Mistral AI·Neon·LangChain으로 Text-to-SQL 변환 시스템 구축하기

자연어 질의를 SQL 문장으로 변환하는 Text-to-SQL 시스템을 RAG 방식으로 만드는 방법을 다루는 포스트예요. 임베딩과 언어 모델에는 Mistral AI를, 벡터 데이터베이스에는 Neon Postgres를 사용하고, LangChain으로 하나로 묶습니다.

출처: 문서

본문

자연어 질의를 SQL 문장으로 변환하는 것은 대규모 언어 모델(LLM)의 강력한 응용이에요. LLM에게 자연어 프롬프트로 SQL을 직접 생성하라고 하는 것도 가능하지만, 한계가 있습니다.

LLM은 SQL 방언(dialect)이 관계형 DB마다 다르기 때문에 문법적으로 올바르지 않은 SQL을 생성할 수 있어요.
LLM은 전체 데이터베이스 스키마·테이블·컬럼 이름·인덱스에 접근할 수 없어 정확하고 효율적인 쿼리를 생성하는 능력이 제한됩니다. 매번 전체 스키마를 입력으로 넣는 것은 비쌉니다.
사전 학습된 LLM은 사용자 피드백과 진화하는 텍스트 질의에 적응할 수 없습니다.

대안은 여러분의 특정 Text-to-SQL 데이터셋(데이터베이스 쿼리 로그와 기타 관련 문맥 포함)으로 LLM을 파인튜닝하는 거예요. 이 접근은 LLM이 정확한 SQL 쿼리를 생성하는 능력을 개선할 수 있지만, 지속적으로 적응하는 데는 여전히 한계가 있습니다. 파인튜닝은 비쌀 수 있어 모델 업데이트 빈도를 제한할 수도 있죠.

LLM은 문맥 내 학습(in-context learning)에 뛰어나므로, 관련 정보를 프롬프트에 넣어 주면 출력을 개선할 수 있어요. 이것이 검색 증강 생성(RAG) 시스템의 아이디어입니다. RAG는 정보 검색과 LLM을 결합해 더 유익하고 문맥적인 응답을 생성하게 합니다.

지식 베이스(데이터베이스 스키마, 질의할 테이블, 이전에 생성된 SQL 쿼리)에서 관련 정보를 검색함으로써 LLM을 활용해 더 정확하고 효율적인 SQL 쿼리를 생성할 수 있어요.

설정과 의존성

Mistral AI API

Mistral AI에 가입하고 콘솔에서 API 키 섹션으로 이동해 새 API 키를 만듭니다. 이 키로 Mistral AI의 임베딩·언어 모델에 접근합니다.

MISTRAL_API_KEY = "your-mistral-api-key"

Neon 데이터베이스

Neon에 가입합니다. Neon 프로젝트에는 neondb라는 바로 사용 가능한 Postgres 데이터베이스가 딸려 있어요. Neon 콘솔의 Connection Details 섹션에서 데이터베이스 연결 문자열을 찾을 수 있습니다. 대략 이런 형태입니다.

postgres://alex:***@ep-cool-darkness-123456.us-east-2.aws.neon.tech/dbname?sslmode=require

NEON_CONNECTION_STRING = "your-neon-connection-string"

Python 라이브러리

RAG 시스템을 만드는 데 필요한 라이브러리를 설치합니다.

%pip install langchain langchain-mistralai langchain-postgres

langchain-postgres는 pgvector가 설치된 Postgres DB에 임베딩을 저장·질의할 수 있게 해 주는 vectorstore 모듈을 제공합니다. 한편 langchain-mistralai는 Mistral 모델과 상호작용하는 데 필요해요.

이 예제에서는 널리 쓰이는 Northwind 샘플 데이터셋을 활용합니다. 고객에게 제품을 판매하는 가상의 무역 회사 Northwind Traders를 모델링하며, Customers·Orders·Products·Employees 같은 테이블이 관계로 연결되어 있어 판매·재고·기타 비즈니스 운영 관련 데이터를 질의·분석할 수 있어요.

Mistral LLM을 호출할 때 두 가지 정보를 문맥으로 제공하려고 합니다.

Northwind 데이터베이스 스키마의 관련 테이블/인덱스 정보
LLM이 학습할 (텍스트 질문, SQL 쿼리) 쌍 몇 개

벡터 데이터베이스를 활용해 스키마와 (텍스트, SQL) 샘플 쌍을 저장하는 검색을 설정합니다. 각 정보 조각에 mistral-embed LLM 모델로 임베딩을 만들고, 쿼리 시점에 질의 임베딩과 저장된 임베딩을 비교해 관련 스니펫을 검색합니다.

데이터 준비

langchain-postgres로 데이터베이스에 데이터의 임베딩을 저장합니다.

import sqlalchemy

# Connect to the database
engine = sqlalchemy.create_engine(
    url=NEON_CONNECTION_STRING, pool_pre_ping=True, pool_recycle=300
)

from langchain_mistralai.embeddings import MistralAIEmbeddings
from langchain_postgres.vectorstores import PGVector
from langchain_core.documents import Document

embeds_model = MistralAIEmbeddings(model="mistral-embed", api_key=MISTRAL_API_KEY)

vector_store = PGVector(
    embeddings=embeds_model,
    connection=engine,
    use_jsonb=True,
    collection_name="text-to-sql-context",
)

다음으로 Northwind 스키마와 샘플 쿼리의 임베딩을 생성합니다. 여기서 PGVector 같은 langchain 벡터 스토어의 add_documents 메서드는 지정된 임베딩 모델로 입력 텍스트의 임베딩을 생성해 데이터베이스에 저장합니다.

참고: Colab에서 작업한다면 데이터베이스 설정과 샘플 쿼리 스크립트를 내려받아 실행하세요.

# import os
# import requests
# repo_url = "https://raw.githubusercontent.com/neondatabase/mistral-neon-text-to-sql/main/data/"
# fnames = ["northwind-schema.sql", "northwind-queries.jsonl"]
# os.mkdir("data")
# for fname in fnames:
#     response = requests.get(repo_url + fname)
#     with open(f"data/{fname}", "w") as file:
#         file.write(response.text)

Northwind 데이터베이스를 만드는 DDL 문장입니다.

# DDL statements to create the Northwind database
_all_stmts = []
with open("data/northwind-schema.sql", "r") as f:
    stmt = ""
    for line in f:
        if line.strip() == "" or line.startswith("--"):
            continue
        else:
            stmt += line
        if ";" in stmt:
            _all_stmts.append(stmt.strip())
            stmt = ""

ddl_stmts = [x for x in _all_stmts if x.startswith(("CREATE", "ALTER"))]

docs = [
    Document(page_content=stmt, metadata={"id": f"ddl-{i}", "topic": "ddl"})
    for i, stmt in enumerate(ddl_stmts)
]
vector_store.add_documents(docs, ids=[doc.metadata["id"] for doc in docs])

샘플 질문-쿼리 쌍입니다.

# Sample question-query pairs
with open("data/northwind-queries.jsonl", "r") as f:
    docs = [
        Document(
            page_content=pair,
            metadata={"id": f"query-{i}", "topic": "query"},
        )
        for i, pair in enumerate(f)
    ]
vector_store.add_documents(docs, ids=[doc.metadata["id"] for doc in docs])

Neon 데이터베이스에 Northwind 테이블도 만들어 LLM 출력을 실행하고 자연어-쿼리 결과 파이프라인이 동작하게 합니다.

# run the DDL script to create the database
with engine.connect() as conn:
    with open("data/northwind-schema.sql") as f:
        conn.execute(sqlalchemy.text(f.read()))
        conn.commit()

SQL 출력 생성하기

마지막으로 Mistral AI의 채팅 모델로 검색된 문맥을 기반으로 SQL 문장을 생성합니다. 먼저 Mistral AI 모델에 넘길 프롬프트를 구성합니다. 프롬프트에는 질의, 검색된 스키마 스니펫, 코퍼스의 유사 쿼리 몇 개가 들어갑니다.

import json

prompt = """
You are an AI assistant that converts natural language questions into SQL queries. To do this, you will be provided with three key pieces of information:
1. Some DDL statements describing tables, columns and indexes in the database:
<schema>
{SCHEMA}
</schema>
2. Some example pairs demonstrating how to convert natural language text into a corresponding SQL query for this schema:
<examples>
{EXAMPLES}
</examples>
3. The actual natural language question to convert into an SQL query:
<question>
{QUESTION}
</question>
Follow the instructions below:
1. Your task is to generate an SQL query that will retrieve the data needed to answer the question, based on the database schema.
2. First, carefully study the provided schema and examples to understand the structure of the database and how the examples map natural language to SQL for this schema.
3. Your answer should have two parts:
- Inside <scratchpad> XML tag, write out step-by-step reasoning to explain how you are generating the query based on the schema, example, and question.
- Then, inside <sql> XML tag, output your generated SQL.
"""

schema = ""
for stmt in relevant_ddl_stmts:
    schema += stmt.page_content + "\n\n"

examples = ""
for stmt in similar_queries:
    text_sql_pair = json.loads(stmt.page_content)
    examples += "Question: " + text_sql_pair["question"] + "\n"
    examples += "SQL: " + text_sql_pair["query"] + "\n\n"

LLM이 단계별로 생각하도록 프롬프트하면 생성 출력의 품질이 좋아집니다. 그래서 LLM에게 추론과 SQL 쿼리를 출력 텍스트의 별도 블록으로 내보내도록 지시합니다.

import re
from langchain_mistralai.chat_models import ChatMistralAI
from langchain_core.messages import HumanMessage

chat_model = ChatMistralAI(api_key=MISTRAL_API_KEY)
response = chat_model.invoke(
    [
        HumanMessage(
            content=prompt.format(QUESTION=question, SCHEMA=schema, EXAMPLES=examples)
        )
    ]
)

sql_query = re.search(r"<sql>(.*?)</sql>", response.content, re.DOTALL).group(1)
print(sql_query)

Mistral AI 모델 출력에서 SQL 문장을 추출해 Neon 데이터베이스에서 실행해 유효한지 확인합니다.

from sqlalchemy import text

with engine.connect() as conn:
    result = conn.execute(text(sql_query))
    for row in result:
        print(row._mapping)

이렇게 Mistral AI API를 채팅·임베딩 모델에, Neon을 벡터 데이터베이스로 활용해 자연어-질의-to-SQL 쿼리 시스템이 동작하게 됐습니다.

결론 (Conclusion)

프로덕션에서 사용하려면 몇 가지 고려 사항을 기억하세요.

생성된 SQL 쿼리를 실행 전에 검증하세요. 특히 DELETE·UPDATE 같은 파괴적 연산은 조심해야 해요. 텍스트 입력이 사용자로부터 오므로 악의적인 입력으로 SQL 인젝션 공격이 일어날 수도 있습니다.
시간에 따라 시스템의 성능·정확도를 모니터링하세요. 데이터가 진화함에 따라 사용 LLM 모델과 지식 베이스 임베딩을 업데이트해야 할 수 있어요.
더 나은 메타데이터. 유사한 예시와 스키마도 유용하지만, 데이터 계보(data lineage)나 대시보드 로그 같은 정보가 LLM API 호출에 더 많은 문맥을 더할 수 있습니다.
검색 개선. 복잡한 질의의 경우 LLM 모델에 넘기는 스키마 정보를 늘려야 할 수 있어요. 또한 우리의 유사도 검색 휴리스틱은 텍스트 질의를 SQL 문장에 매칭하는 단순한 방식입니다. HyDE(Hypothetical Document Expansion) 같은 기법을 쓰면 검색 스니펫의 품질을 개선할 수 있습니다.

부록 (Appendix)

Northwind 데이터베이스 설정 스크립트와 샘플 쿼리는 다음 저장소에서 가져왔어요.

Northwind Psql
Sample queries

더 알아보기 (Learn more)

Neon 공식 문서 — 서버리스 PostgreSQL 데이터베이스
pgvector 문서 — Postgres 벡터 확장
langchain_mistralai의 MistralAIEmbeddings / ChatMistralAI — Mistral 통합
mistral-embed — 텍스트 임베딩용 모델

Mistral AI·Neon·LangChain으로 Text-to-SQL 변환 시스템 구축하기