Vanna와 Milvus를 활용한 자연어 SQL 쿼리 시스템 구축하기

이번 포스트에서는 Vanna와 Milvus를 사용하여 자연어로 SQL 쿼리를 생성하는 시스템을 구축하는 방법을 살펴보겠습니다.

출처: https://vanna.ai/

1. 시스템 구성 요소

사용된 주요 라이브러리

• Vanna: 자연어를 SQL로 변환해주는 프레임워크
• Milvus: 벡터 데이터베이스
• OpenAI: 언어 모델 기반 처리
• SQLite: 관계형 데이터베이스

2. 코드 구조 분석

환경 설정

import os
import sqlite3
from pymilvus import MilvusClient, model
from vanna.milvus import Milvus_VectorStore
from vanna.openai import OpenAI_Chat

Vanna-Milvus 클래스 정의

class VannaMilvus(Milvus_VectorStore, OpenAI_Chat):
    def __init__(self, config=None):
        Milvus_VectorStore.__init__(self, config=config)
        OpenAI_Chat.__init__(self, config=config)

Milvus_VectorStore와 OpenAI_Chat을 상속받아 벡터 저장소와 챗봇 기능을 통합했습니다.

Milvus 클라이언트 설정

milvus_uri = "./milvus_vanna.db"
milvus_client = MilvusClient(uri=milvus_uri)
 
vn_milvus = VannaMilvus(
    config = {
        "openai_base_url": "http://localhost:1234/v1",
        "api_key": os.getenv("OPENAI_API_KEY"),
        "model": "qwen2.5-14b-instruct",
        "milvus_client": milvus_client,
        "embedding_function": model.DefaultEmbeddingFunction(),
        "n_results": 2,
    }
)

3. 데이터베이스 구조

테이블 구성

• Customer: 고객 정보
• Company: 회사 정보
• User: 사용자 정보

각 테이블은 기본적인 CRUD 작업을 위한 필드들을 포함하고 있습니다.

4. 학습 데이터 관리

기존 학습 데이터 제거

existing_training_data = vn_milvus.get_training_data()
if len(existing_training_data) > 0:
    for _, training_data in existing_training_data.iterrows():
        vn_milvus.remove_training_data(training_data["id"])

새로운 학습 데이터 추가

1. DDL 정보 학습
2. 비즈니스 용어 정의 추가
3. 예시 SQL 쿼리 추가

5. 시스템의 주요 특징

1. 벡터 기반 검색: Milvus를 통한 효율적인 시맨틱 검색
2. 자연어 처리: OpenAI 모델을 활용한 자연어 이해
3. 확장 가능한 구조: 새로운 학습 데이터 쉽게 추가 가능
4. 유연한 설정: 다양한 설정 옵션 제공

결론

이 시스템은 자연어로 데이터베이스를 쿼리할 수 있게 해주는 강력한 도구입니다. Milvus의 벡터 검색 능력과 OpenAI의 언어 모델을 결합하여 효율적이고 직관적인 데이터베이스 접근을 가능하게 합니다.

더 자세한 내용이 필요하신가요?

 

전체 코드 파일을 아래 첨부합니다.

vannaAIwithMilvus_src.zip

0.00MB