[9편] AI 모델과 pgvector 연동: 임베딩 기반 검색 시스템 구축

 

이 글에서는 AI 모델을 사용하여 문장을 벡터(Embedding)로 변환하고, pgvector를 활용하여 AI 검색 시스템을 구축하는 방법을 다룹니다.

✅ OpenAI text-embedding-ada-002 및 Hugging Face sentence-transformers 활용

✅ 임베딩 벡터를 pgvector에 저장하고, 유사한 검색 결과 반환

✅ FastAPI를 이용해 벡터 검색 API 구축

 

---

🚀 1. AI 모델을 활용한 벡터 변환(Embedding)

 

1️⃣ AI 임베딩 모델이란?

 

임베딩(Embedding)은 텍스트, 이미지, 음성 등의 데이터를 벡터 형태로 변환하는 기술입니다.

이 벡터는 유사한 의미를 가지는 데이터끼리 가까운 위치에 배치되도록 학습됩니다.

 

✅ AI 검색 시스템에서 임베딩 활용 방식

1️⃣ 문서 또는 텍스트 데이터를 벡터로 변환

2️⃣ 변환된 벡터를 PostgreSQL(pgvector)에 저장

3️⃣ 사용자의 입력(쿼리)을 벡터로 변환 후 가장 유사한 데이터를 검색

 

✅ 대표적인 임베딩 모델

• OpenAI text-embedding-ada-002 → 1536차원 벡터 출력 (OpenAI API 필요)

• Hugging Face sentence-transformers → 768차원 벡터 출력 (오픈소스, 로컬 실행 가능)

 

---

2️⃣ OpenAI text-embedding-ada-002 활용

 

OpenAI API를 사용하여 문장을 벡터로 변환할 수 있습니다.

 

📌 Python 라이브러리 설치

pip install openai

 

📌 Python 코드: 문장을 벡터로 변환

import openai

openai.api_key = "your-openai-api-key"

def get_openai_embedding(text: str):
    response = openai.Embedding.create(
        input=text,
        model="text-embedding-ada-002"
    )
    return response['data'][0]['embedding']

# 테스트 실행
vector = get_openai_embedding("PostgreSQL is an advanced database.")
print(len(vector))  # 1536차원 벡터 출력

 

✅ 출력 예시

1536

 

 

---

3️⃣ Hugging Face sentence-transformers 활용 (로컬 실행)

 

Hugging Face의 sentence-transformers는 로컬에서 실행 가능하여 API 비용 없이 벡터 변환이 가능합니다.

 

📌 Python 라이브러리 설치

pip install sentence-transformers

 

📌 Python 코드: 문장을 벡터로 변환

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

def get_huggingface_embedding(text: str):
    return model.encode(text).tolist()

# 테스트 실행
vector = get_huggingface_embedding("PostgreSQL is an advanced database.")
print(len(vector))  # 768차원 벡터 출력

 

✅ 출력 예시

768

 

 

---

🚀 2. pgvector 테이블 생성 및 벡터 저장

 

임베딩 벡터를 저장할 pgvector 테이블을 PostgreSQL에 생성합니다.

 

📌 PostgreSQL 접속

kubectl exec -it $(kubectl get pod -n database -l app.kubernetes.io/name=postgresql -o jsonpath="{.items[0].metadata.name}") -n database -- psql -U postgres -d ragdb

 

📌 벡터 저장 테이블 생성 (1536차원)

CREATE TABLE IF NOT EXISTS embeddings (
    id SERIAL PRIMARY KEY,
    content TEXT,
    embedding vector(1536)
);

 

✅ 테이블 생성 확인

\dt

 

✅ 출력 예시

          List of relations
 Schema |    Name    | Type  |  Owner  
--------+-----------+-------+---------
 public | embeddings | table | postgres

 

 

---

3️⃣ 벡터 데이터 삽입

 

Python 코드에서 pgvector에 벡터를 저장하는 API를 만듭니다.

 

📌 Python 코드: 벡터 저장 API (FastAPI 사용)

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import psycopg2
import openai
from sentence_transformers import SentenceTransformer
from typing import List

app = FastAPI()

# PostgreSQL 연결 정보
DB_HOST = "localhost"
DB_PORT = "5432"
DB_NAME = "ragdb"
DB_USER = "postgres"
DB_PASSWORD = "postgresql"

# OpenAI API 설정
openai.api_key = "your-openai-api-key"

# Hugging Face 모델 로드
hf_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# PostgreSQL 연결 함수
def get_db_connection():
    return psycopg2.connect(
        host=DB_HOST, port=DB_PORT,
        database=DB_NAME, user=DB_USER, password=DB_PASSWORD
    )

# 데이터 모델 정의
class TextData(BaseModel):
    content: str
    use_openai: bool = True  # True면 OpenAI, False면 Hugging Face 사용

# 벡터 삽입 API
@app.post("/add_text/")
def add_text(data: TextData):
    conn = get_db_connection()
    cur = conn.cursor()

    # 벡터 생성 (OpenAI 또는 Hugging Face 선택)
    if data.use_openai:
        response = openai.Embedding.create(input=data.content, model="text-embedding-ada-002")
        embedding = response['data'][0]['embedding']
    else:
        embedding = hf_model.encode(data.content).tolist()

    try:
        # 벡터를 PostgreSQL에 삽입
        cur.execute(
            "INSERT INTO embeddings (content, embedding) VALUES (%s, %s) RETURNING id;",
            (data.content, embedding)
        )
        conn.commit()
        return {"message": "Vector added successfully", "id": cur.fetchone()[0]}
    
    except Exception as e:
        conn.rollback()
        raise HTTPException(status_code=500, detail=str(e))
    
    finally:
        cur.close()
        conn.close()

 

✅ FastAPI 실행

uvicorn main:app --host 0.0.0.0 --port 8000 --reload

 

✅ API 테스트 (http://localhost:8000/docs)

• POST /add_text/ 를 사용하여 문장을 벡터로 변환 & 저장

 

---

🚀 3. 벡터 검색 API 구현

 

FastAPI를 활용하여 pgvector에서 벡터 데이터를 검색하는 API를 추가합니다.

 

📌 Python 코드: 벡터 검색 API

from fastapi import Query

# 검색 API
@app.get("/search/")
def search_vector(query_text: str, use_openai: bool = True):
    conn = get_db_connection()
    cur = conn.cursor()

    # 검색 벡터 생성
    if use_openai:
        embedding = get_openai_embedding(query_text)
    else:
        embedding = get_huggingface_embedding(query_text)

    try:
        cur.execute(
            "SELECT content, embedding <-> %s AS similarity FROM embeddings ORDER BY similarity LIMIT 5;",
            (embedding,)
        )
        results = cur.fetchall()
        return [{"content": row[0], "similarity": row[1]} for row in results]

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    
    finally:
        cur.close()
        conn.close()

 

✅ API 테스트 (http://localhost:8000/docs)

• GET /search/?query_text=What is PostgreSQL?

• 가장 유사한 문장 5개 반환

 

---

📌 4. 최종 정리

 

✅ OpenAI & Hugging Face 임베딩 모델 활용

✅ pgvector 테이블 생성 및 벡터 데이터 저장

✅ FastAPI 기반 벡터 검색 API 구축