메달리온 데이터 위에 Text-to-SQL과 RAG 올리기

# 개념 예시 — 실제 프롬프트는 도메인에 맞게 조정해야 한다
ROUTER_PROMPT = """Classify the user's question into exactly one category.
- "data": metrics, counts, trends, statistics, numbers
- "docs": documentation, processes, architecture, how-to, team info
Respond with ONLY the word "data" or "docs"."""

# 개념 예시 — 스키마 인덱스에 넣을 문서 구조
schema_documents = [
    {
        "doc_id": "daily_event_summary",
        "table_name": "daily_event_summary",
        "content": "일별 이벤트 집계 테이블. 이벤트 유형별 건수, 고유 사용자 수를 보여준다.",
        "columns_json": '{"event_date": "DATE", "event_type": "STRING", "event_count": "BIGINT", ...}',
        "sample_queries": "SELECT event_type, SUM(event_count) FROM ... GROUP BY 1"
    },
    # ... Gold 테이블마다 하나씩
]

# 개념 예시 — 자연어 질문과 검증된 SQL 쌍
example_documents = [
    {
        "question": "How many events happened last month?",
        "question_ko": "지난 달 이벤트가 몇 건이었나?",
        "sql_query": """SELECT event_type, COUNT(*) AS cnt
FROM daily_event_summary
WHERE event_date >= date_trunc('MONTH', current_date() - INTERVAL 1 MONTH)
  AND event_date < date_trunc('MONTH', current_date())
GROUP BY 1""",
        "explanation": "일별 집계 테이블에서 지난 달 이벤트를 유형별로 합산"
    },
    # ... 수십 개의 예시
]

# 개념 예시 — Text-to-SQL 파이프라인 흐름
def generate_sql(question: str) -> str:
    # 1. 질문과 관련된 테이블 스키마 검색
    schemas = schema_index.similarity_search(question, num_results=5)
 
    # 2. 유사 질문에 대한 검증된 SQL 예시 검색
    examples = example_index.similarity_search(question, num_results=3)
 
    # 3. 검색 결과를 프롬프트에 조합
    prompt = f"""You are a SQL expert.
 
Available tables:
{format_schemas(schemas)}
 
Similar examples:
{format_examples(examples)}
 
Current date: {current_date()}
 
Generate a SELECT query for: {question}
Only SELECT is allowed. Always add LIMIT."""
 
    # 4. LLM이 SQL 생성
    sql = llm.generate(prompt)
 
    # 5. 안전 검증 (SELECT만 허용, DDL/DML 차단)
    validate_sql(sql)
 
    return sql

# 개념 예시 — RAG 검색 및 답변 생성 흐름
def rag_answer(question: str):
    # 1. 질문과 유사한 문서 청크 검색
    results = vector_index.similarity_search(
        query_text=question,
        columns=["text", "url", "page_title"],
        num_results=5
    )
 
    # 2. 검색된 문서를 LLM 컨텍스트로 구성
    context = "\n".join([
        f"[Document {i}] {r.page_title}\n{r.text}"
        for i, r in enumerate(results)
    ])
 
    # 3. LLM이 컨텍스트 기반으로 답변 생성
    answer = llm.generate(
        f"Based on the following documents, answer: {question}\n\n{context}"
    )
 
    # 답변과 출처 URL을 함께 반환
    return answer, [r.url for r in results]