PostgreSQL과 Milvus를 동일한 ID로 상호 연결하여 RAG를 구현

PostgreSQL과 Milvus를 동일한 ID로 상호 연결하여 RAG를 구현하는 간단한 Python 예제 코드를 작성해줘

PostgreSQL(RDBMS)에는 원문 텍스트와 메타데이터를 저장하고, Milvus(Vector DB)에는 임베딩 벡터를 저장한 뒤 동일한 doc_id로 상호 참조하여 RAG(검색 증강 생성)를 수행하는 파이썬 코드입니다.

이 구조를 이용하면 Milvus로 유사한 문장을 빠르게 찾은 후, PostgreSQL에서 해당 문장의 정확한 원문과 작성자/생성일 등 메타데이터를 한 번에 조회할 수 있습니다.

1. 사전 라이브러리 설치

WSL2 우분투 환경 또는 파이썬 가상환경에서 필요한 라이브러리를 설치합니다.

Bash

pip install pymilvus psycopg2-binary ollama

2. 파이썬 연동 전체 예제 코드 (rag_example.py)

Python

import uuid

import psycopg2

from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility

import ollama

 

# ==========================================

# 1. 데이터베이스 연결 설정

# ==========================================

# PostgreSQL 연결

pg_conn = psycopg2.connect(

    host="localhost",

    port=5432,

    dbname="aidb",

    user="myuser",

    password="mypassword"

)

pg_cursor = pg_conn.cursor()

 

# PostgreSQL 테이블 생성 (원문 및 메타데이터 저장용)

pg_cursor.execute("""

CREATE TABLE IF NOT EXISTS documents (

    doc_id VARCHAR(64) PRIMARY KEY,

    title VARCHAR(255),

    content TEXT,

    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP

);

""")

pg_conn.commit()

 

# Milvus 연결

connections.connect("default", host="localhost", port="19530")

 

# Milvus 컬렉션 생성 (벡터 저장용)

collection_name = "doc_vectors"

 

# 이미 컬렉션이 있다면 로드, 없으면 새로 정의

if utility.has_collection(collection_name):

    milvus_collection = Collection(collection_name)

else:

    fields = [

        # PostgreSQL의 doc_id와 동일하게 사용할 문자열 PK

        FieldSchema(name="doc_id", dtype=DataType.VARCHAR, is_primary=True, max_length=64),

        # Ollama 'nomic-embed-text' 기준 벡터 차원수 (768)

        FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=768)

    ]

    schema = CollectionSchema(fields, description="RAG Vector Store")

    milvus_collection = Collection(collection_name, schema)

   

    # 인덱스 생성 (빠른 유사도 검색용)

    index_params = {

        "metric_type": "COSINE",

        "index_type": "IVF_FLAT",

        "params": {"nlist": 128}

    }

    milvus_collection.create_index(field_name="vector", index_params=index_params)

 

milvus_collection.load()

 

# ==========================================

# 2. 문서 저장 함수 (PostgreSQL + Milvus 동시 저장)

# ==========================================

def insert_document(title: str, content: str):

    # 1. 두 DB가 함께 사용할 고유 ID 생성

    doc_id = str(uuid.uuid4())

 

    # 2. Ollama를 이용해 텍스트를 벡터로 변환

    embed_response = ollama.embeddings(model="nomic-embed-text", prompt=content)

    embedding_vector = embed_response["embedding"]

 

    # 3. PostgreSQL에 원문 저장

    pg_cursor.execute(

        "INSERT INTO documents (doc_id, title, content) VALUES (%s, %s, %s)",

        (doc_id, title, content)

    )

    pg_conn.commit()

 

    # 4. Milvus에 벡터 및 doc_id 저장

    milvus_collection.insert([

        [doc_id],

        [embedding_vector]

    ])

    milvus_collection.flush()

    print(f"✅ 문서 저장 완료 [ID: {doc_id}] - 제목: {title}")

 

# ==========================================

# 3. RAG 검색 및 답변 생성 함수

# ==========================================

def ask_question(question: str):

    print(f"\n❓ 질문: {question}")

   

    # 1. 질문을 벡터로 변환

    q_embed = ollama.embeddings(model="nomic-embed-text", prompt=question)["embedding"]

 

    # 2. Milvus에서 벡터 유사도 상위 1개 검색

    search_params = {"metric_type": "COSINE", "params": {"nprobe": 10}}

    results = milvus_collection.search(

        data=[q_embed],

        anns_field="vector",

        param=search_params,

        limit=1,

        output_fields=["doc_id"]

    )

 

    if not results or len(results[0]) == 0:

        print("관련 문서를 찾지 못했습니다.")

        return

 

    # 3. 매칭된 Milvus의 doc_id 추출

    matched_doc_id = results[0][0].id

    score = results[0][0].distance

    print(f"🔍 Milvus에서 검색된 doc_id: {matched_doc_id} (유사도 스코어: {score:.4f})")

 

    # 4. 해당 doc_id로 PostgreSQL에서 원문 가져오기 (상호 참조)

    pg_cursor.execute(

        "SELECT title, content, created_at FROM documents WHERE doc_id = %s",

        (matched_doc_id,)

    )

    doc = pg_cursor.fetchone()

   

    if doc:

        title, content, created_at = doc

        print(f"📖 PostgreSQL 조회 결과 -> [제목]: {title} | [작성일]: {created_at}")

 

        # 5. Qwen2.5에 원문(Context)과 질문 전달하여 답변 생성

        prompt = f"""다음 바탕 지식을 참고하여 질문에 친절하게 답변해 주세요.

 

[바탕 지식]

제목: {title}

내용: {content}

 

질문: {question}

답변:"""

 

        response = ollama.generate(model="qwen2.5", prompt=prompt)

        print("\n🤖 Qwen의 답변:")

        print(response["response"])

 

# ==========================================

# 4. 실행 테스트

# ==========================================

if __name__ == "__main__":

    # 임베딩 전용 모델 설치 확인 (최초 1회 필요)

    # 터미널 명령어: ollama pull nomic-embed-text

 

    # 샘플 문서 저장

    insert_document(

        title="동아리 동호회 규칙",

        content="컴퓨터 동아리의 정기 모임은 매주 목요일 저녁 7시 동아리방에서 열립니다. 지각 시 회비 1000원이 부과됩니다."

    )

    insert_document(

        title="학교 도서관 이용 안내",

        content="도서관 대출 가능 권수는 1인당 최대 5권이며, 대출 기간은 주말 포함 14일입니다."

    )

 

    # RAG 질문 테스트

    ask_question("동아리 모임에 늦으면 벌금이 얼마야?")

3. 동작 방식 요약

1.    insert_document():

o   uuid로 공통 고유값 doc_id를 하나 만듭니다.

o   nomic-embed-text 모델로 문장을 벡터로 바꾼 후, Milvus에는 (doc_id, vector)를 저장합니다.

o   PostgreSQL에는 (doc_id, title, content, created_at) 원본 데이터를 저장합니다.

2.    ask_question():

o   사용자의 질문을 벡터로 변환해 Milvus에서 제일 유사한 벡터의 doc_id를 찾습니다.

o   그 doc_id를 조건(WHERE doc_id = %s)으로 PostgreSQL에 SQL Query를 날려 정확한 원문 텍스트를 조회합니다.

o   가져온 원문을 Qwen2.5 모델에게 참고 자료로 제공하여 최종 답변을 받습니다.

임베딩 모델(nomic-embed-text) 다운로드 및 Ollama 연동 방법이 궁금하신가요?

예

Ollama에서 nomic-embed-text 임베딩 모델을 다운로드하고 파이썬과 연동하는 법을 알려줘

nomic-embed-text는 문장을 숫자로 바꾸는(임베딩) 속도가 빠르고 성능이 뛰어난 대표적인 오픈소스 임베딩 모델입니다.

Ollama에서 이 모델을 다운로드하고 파이썬 코드로 불러와 다루는 방법을 가장 쉽게 정리해 드립니다.

1. Ollama에서 nomic-embed-text 모델 다운로드

WSL2 우분투 터미널을 열고 다음 명령어를 입력하여 임베딩 전용 모델을 받습니다.

Bash

# nomic-embed-text 모델 받기 (약 274MB)

ollama pull nomic-embed-text

다운로드가 완료된 후, 모델이 잘 설치되었는지 확인합니다.

Bash

ollama list

목록에 nomic-embed-text:latest가 보이면 준비가 완료된 것입니다.

2. 파이썬 환경 설정

파이썬 환경에서 Ollama를 제어하기 위한 라이브러리를 설치합니다.

Bash

pip install ollama numpy

3. 파이썬 연동 기본 코드

파이썬 파일(예: embed_test.py)을 생성하고 아래 코드를 실행해 봅니다.

Python

import ollama

import numpy as np

 

# 1. 단일 문장 임베딩 생성

text = "인공지능과 데이터베이스를 연동하는 방법"

 

response = ollama.embeddings(

    model="nomic-embed-text",

    prompt=text

)

 

# 생성된 벡터 추출 (768차원의 실수 리스트)

vector = response["embedding"]

 

print(f"✅ 임베딩 완료!")

print(f"- 문장: {text}")

print(f"- 벡터 차원 수: {len(vector)}")  # 결과: 768

print(f"- 벡터 샘플 (앞 5개 숫자): {vector[:5]}")

 

 

# 2. 두 문장의 유사도(Cos Similarity) 비교 예시

def cosine_similarity(v1, v2):

    return np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))

 

text1 = "학교 도서관은 저녁 9시까지 운영합니다."

text2 = "도서관 문 닫는 시간이 몇 시인가요?"

text3 = "오늘 점심 메뉴는 돈가스입니다."

 

# 각 문장의 벡터 구하기

v1 = ollama.embeddings(model="nomic-embed-text", prompt=text1)["embedding"]

v2 = ollama.embeddings(model="nomic-embed-text", prompt=text2)["embedding"]

v3 = ollama.embeddings(model="nomic-embed-text", prompt=text3)["embedding"]

 

print("\n📊 문장 의미 유사도 측정 (1에 가까울수록 비슷한 의미):")

print(f"- [문장1 vs 문장2 (관련 질문)]: {cosine_similarity(v1, v2):.4f}")

print(f"- [문장1 vs 문장3 (무관한 주제)]: {cosine_similarity(v1, v3):.4f}")

4. 핵심 포인트 요약

• 차원 수 (Dimension): nomic-embed-text는 결과값으로 768개의 숫자 배열을 반환합니다. 따라서 Milvus의 컬렉션 필드 설정 시 dim=768로 맞춰야 합니다.
• 임베딩 검색 원리: 위 예제처럼 의미가 비슷한 문장(문장1과 문장2)은 벡터 간 유사도 점수가 높게 나오며, Milvus는 이 숫자를 이용해 가장 연관성 있는 원문을 찾아냅니다.