PGVector와 Python FastAPI를 연동하여 벡터 데이터 저장 및 유사도 기반 조회하기

포스팅 개요

이번 포스팅은 PostgreSQL의 PGVector extension을 활용해 벡터 데이터베이스로 사용하여 파이썬(Python)의 FastAPI를 연동해 데이터를 저장하고 조회하는 방법에 대해 정리하는 포스팅입니다. 이때, PostgreSQL에 데이터를 저장하는 방법에는 벡터 데이터베이스로 활용하므로 일반 데이터를 저장하면서 동시에 임베딩 모델(embedding model)을 활용해 텍스트를 벡터(vector)로 변환하여 저장하게 됩니다. 또한, 데이터를 조회하는 과정은 1) 제목(title)과 완벽하게 일치하는 exact match 기반 검색과 2) 코사인 유사도(cosine similarity) 기반으로 텍스트 벡터 유사도 기반으로 검색을 하는 과정을 정리합니다.

PostgreSQL와 PGVecotor란 무엇인지 궁금하시거나 설치하는 방법이 궁금하시다면, 앞서 제가 작성한 PostgreSQL PGVector 설치 및 사용하기 포스팅(https://lsjsj92.tistory.com/675)을 참고해주세요.

PostgreSQL PGVector 설치 및 사용하기(Feat. 벡터 데이터베이스(Vector Database) 구축)

 

또한, 본 포스팅을 작성하기 위해서 참고한 자료는 아래와 같습니다.

• https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2
• https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search/cross_encoder.py
• https://github.com/pgvector/pgvector-python

GitHub - pgvector/pgvector-python: pgvector support for Python

 

본 포스팅에서 사용한 코드는 아래 github에 올려두었으니 참고해주세요.

• https://github.com/lsjsj92/pgvector-py-tutorial

GitHub - lsjsj92/pgvector-py-tutorial: pgvector tutorial with python

 

---

포스팅 본문

포스팅 개요에서도 언급하였듯, 이번 포스팅은 파이썬(Python)의 FastAPI를 활용해서 PGVector가 확장된 PostgreSQL과(Vector database) 연동하여 데이터를 저장 및 조회하는 과정을 정리하는 포스팅입니다. 본 포스팅은 아래와 같은 단계로 정리를 진행하겠습니다.

 

1. 임베딩 모델 준비 - 허깅페이스(HuggingFace) 임베딩 모델 활용

2. FastAPI 환경 개발 - PostgreSQL과 연동하기

3. FastAPI 환경 개발 - 스키마 및 모델 구조 설정하기

4. FastAPI 환경 개발 - 데이터 저장하기

5. FastAPI 환경 개발 - 데이터 조회하기

 

참고사항: PostgreSQL 및 PGVector 그리고 Python FastAPI를 개발한 환경은 다음과 같습니다.

- MacOS(MacBook pro, 2019)
- PostgreSQL version: PostgreSQL@16
- PGVector version: pgvector 0.8.0
- Python version: Python 3.9
- Python library: langchaib, pydantic, fastapi, SQLAlchemy, huggingface, pgvector

 

---

임베딩 모델 준비

가장 먼저, 임베딩 모델을 준비합니다. PGVector를 활용한다는 것은 벡터 데이터베이스(Vector Database)를 사용하는 것이기에 당연히 벡터로 변환해줄 임베딩 모델이 필요하죠. 저는 텍스트 데이터를 vector로 변환하여 저장할 것이기 때문에, 허깅페이스에서 모델을 선정하였습니다. 제가 선정한 모델은 다음과 같습니다.

 

 

sentence-transformers/all-MiniLM-L6-v2라는 모델이며, 임베딩 벡터 차원수도 크지 않고(384차원) 준수한 성능을 보여주기에 해당 모델로 선정하였습니다. 

이 모델을 활용해 텍스트 데이터를 벡터로 변환하는 작업은 아래와 같은 파이썬 코드로 실행할 수 있습니다.

 

hf_embed_repo_id = 'sentence-transformers/all-MiniLM-l6-v2'
hf_embeddings = HuggingFaceHubEmbeddings(repo_id=hf_embed_repo_id)
query_result = hf_embeddings.embed_query("안녕하세요")
print(len(query_result)) # 384 차원

 

추후 FastAPI 환경에도 위와 같은 코드를 기반으로 코드 작업을 진행합니다. 벡터를 저장하는 컬럼에 같은 벡터 차원으로 설정하고(384차원) 이를 활용해 코사인 유사도와 같은 작업을 수행하게 됩니다.

---

FastAPI 개발: PostgreSQL과 Python 연동하기

이제 본격적으로 벡터 데이터베이스(Vector Database)인 PGVector 확장자가 갖추어진 PostgreSQL과 연동하는 파이썬(Python) FastAPI 코드를 살펴보려고 합니다. 

첫 번째로는 Python 환경에서 PostgreSQL과 연동하기 위한 환경 설정을 진행합니다. 아래 코드는 데이터베이스 연결 설정, 세션 생성, 그리고 PGVector 확장을 초기화하는 코드입니다.

 

DB_USER = os.getenv("DB_USER", "leesoojin")
DB_PASS = os.getenv("DB_PASS", "")
DB_HOST = os.getenv("DB_HOST", "localhost")
DB_PORT = os.getenv("DB_PORT", "5432")
DB_NAME = os.getenv("DB_NAME", "test")

 

코드의 첫 번째 부분에는 데이터베이스 연결에 필요한 정보를 환경 변수에서 가져옵니다. 환경 변수를 사용하지 않을 경우 기본 값을 제공하는데요. 데이터베이스 유저와 호스트, 포트, 데이터베이스 이름 등을 작성합니다. 이 정보는 이전 포스팅에서 작성했던 정보를 기반으로 넣어두었습니다.

저는 읽기 쉽게 위와 같이 작성하였지만, .env 파일을 활용해 사용하는 것을 권장드립니다.

 

# PostgreSQL URL (asyncpg)
DATABASE_URL = f"postgresql+asyncpg://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"

# echo=True -> SQL 로그 출력
async_engine = create_async_engine(DATABASE_URL, echo=True)

# 세션 팩토리 (자동커밋/오토플러시 끔)
AsyncSessionLocal = sessionmaker(
    bind=async_engine,
    expire_on_commit=False,
    class_=AsyncSession,
    autoflush=False,
    autocommit=False
)

Base = declarative_base()

 

그리고 설정 정보를 기반으로 데이터베이스 연결을 설정합니다. DATABASE_URL에 관련 정보를 함께 넣어주고 이 정보를 기반으로 데이터베이스 엔진을 생성합니다. 저는 로그를 보기 위해서 echo를 True로 설정하였습니다.

그 다음으로 세션 메이커(sessionmaker)를 이용해 세션 팩토리를 구성하고 declarative_base() 객체를 생성하여 SQLAlchemy의 ORM(Object Relational Mapping)을 사용해 데이터베이스 테이블을 Python 클래스에 매핑하기 위한 준비를 진행합니다.

이 Base라는 객체는 뒤에서 테이블 클래스에서 사용되게 됩니다.

 

async def init_db():
    """
    pgvector 확장을 활성화하고, Base에 정의된 모든 테이블을 생성
    """
    async with async_engine.begin() as conn:
        # pgvector extension
        await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
        # 테이블 생성
        await conn.run_sync(Base.metadata.create_all)

# FastAPI 의존성 주입용
async def get_db_session():
    """
    DB 세션을 제공하는 Generator
    """
    async with AsyncSessionLocal() as session:
        yield session

 

init_db 함수는 PGVector 확장을 활성화하고 위에서 생성한 SQLAlchemy의 Base 객체를 기반으로 모든 테이블을 생성하도록 수행하는 함수이고, get_db_session 함수는 데이터베이스 세션을 제공하는 함수입니다.

즉, 여러분들이 데이터베이스의 테이블을 생성하지 않아도, create_all이라는 것을 활용해 base가 가지고 있는 모든 테이블을 생성하는 것입니다.

---

FastAPI 개발: 데이터 스키마 및 데이터베이스 모델 구조 정의하기

다음으로 PostgreSQL에 PGVector 확장을 활용하여 텍스트 데이터를 저장하고 벡터화된 임베딩을 함께 관리하기 위한 데이터 모델을 정의합니다. 데이터베이스 테이블, 데이터 스키마, 그리고 API 응답 모델을 정의합니다.

저는 먼저, API 응답에 활용되는 2개의 클래스를 구성합니다. 이때 Pydantic을 사용하여 데이터 검증 및 직렬화를 수행합니다. 

from pydantic import BaseModel

class TextItemCreate(BaseModel):
    title: str
    content: str

class TextItemResponse(BaseModel):
    id: int
    title: str
    content: str

    class Config:
        orm_mode = True

 

TextItemCreate 클래스는 텍스트 데이터를 생성하거나 저장 요청을 할 때 사용되는 데이터 스키마입니다. 이때 제목(title)과 콘텐츠(content) 정보를 받도록 합니다.

TextItemResponse 클래스는 데이터 응답을 원할 때 사용되는 데이터 스키마입니다. 즉, 사용자의 요청(request)에 따라 응답(response)를 수행하는데, 이때 id와 title, content가 포함되어 응답하도록 정의된 데이터 스키마이죠.

 

또한, 데이터베이스 테이블에 대한 정보도 설정해두겠습니다. SQLAlchemy는 Python의 ORM 라이브러리인데요. 이를 활용하여 Python 클래스를 데이터베이스 테이블에 매핑하여 DB 작업을 처리할 수 있도록 해줍니다.

class TextItem(Base):
    __tablename__ = "text_items"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    title = Column(String, nullable=False)   # 제목
    content = Column(Text, nullable=False)   # 본문
    embed = Column(Vector(384), nullable=False)  # 384차원 벡터

 

위 클래스에서 TextItem(Base)라고 되어 있는데요. 이때 Base가 앞서 정의했던 declarative_base() 객체입니다. 이 클래스는 데이터베이스의 테이블 구조입니다. 이 테이블에는 id, title, content, embed 4개의 컬럼이 있는 것입니다.

이때, embed가 바로 벡터를 저장하는 컬럼입니다. 즉, 이를 활용해 벡터 데이터베이스(Vector Database) 역할을 수행할 수 있도록 하는 것이죠. 또한, 벡터 차원을 384으로 설정해두었는데 위에서 허깅페이스 임베딩 모델의 차원이 384차원의 모델이기 때문에 이와 같이 설정하였습니다.

반응형

---

FastAPI 개발: 데이터 저장하기

이번에는 FastAPI를 활용하여 받은 데이터를 PostgreSQL에 저장하는 방법을 살펴봅니다. 이 작업은 입력 데이터 중 콘텐츠(content) 데이터를 앞서 정의한 임베딩 모델을 활용해 embedding vector로 변환한 뒤 PGVector를 통해 데이터베이스에 저장하는 과정을 포함합니다. 이렇게 데이터를 저장하면 벡터 데이터베이스로 활용할 준비 작업이 진행되는 것이죠.

@router.post("/create-item", response_model=TextItemResponse, tags=["Items"])
async def create_text_item(
    item: TextItemCreate,
    db: Session = Depends(get_db_session)
):
    """
    입력 텍스트(title, content)를 임베딩 후 PostgreSQL에 저장
    """
    vector = get_embedding(item.content)  # content 임베딩
    db_item = TextItem(
        title=item.title,
        content=item.content,
        embed=vector
    )
    db.add(db_item)
    
    await db.commit()
    await db.refresh(db_item)
    
    return db_item

 

이 API는 title과 content가 포함된 POST 요청을 /create-item 이라는 엔드포인트로 처리합니다. 요청 데이터는 Pydantic을 활용해 정의된 스키마 클래스 TextItemCreate를 기반으로 검증합니다.

검증된 데이터는 먼저 get_embedding 함수를 사용해 content 텍스트 데이터를 임베딩합니다. 이 함수는 허깅페이스(HuggingFace) 임베딩 모델을 활용해 384차원의 벡터로 변환합니다. 이렇게 생성된 벡터는 데이터베이스에 저장될 PGVector 필드(embed)에 저장될 준비를 합니다.

이후 SQLAlchemy ORM 클래스 TextItem를 사용해 데이터베이스 테이블(text_Items)의 행을 정의합니다. 정의한다는 것은,

- title에는 요청 받은 item.title을 값을 저정하고

- content에는 item.content 데이터를

- embed 컬럼에는 방금 생성한 임베딩 벡터 값인 vector를 할당한다는 것을 의미합니다.

 

구성된 데이터는 db.add를 통해 데이터베이스 세션에 추가됩니다. 이후 commit을 호출하여 실제로 저장되게 되고 refresh를 사용해 데이터베이스에서 갱신된 객체를 가져오게 됩니다. 이 과정에서 DB에 자동으로 생성된 값(id)가 반영되죠. 

 

실제로 실행한 결과를 살펴보겠습니다. 아래와 같이 데이터를 직접 넣었을 때 DB에 정상적으로 저장이 될까요?

 

 

FastAPI에서 제공해주는 Swagger 화면에서 데이터를 직접 넣어보았습니다. title과 content에 알맞는 데이터를 넣고 실행해보겠습니다.

 

 

그 결과 API 응답 코드가 200 코드가 나와서 정상적으로 동작되었음을 확인할 수 있습니다.

또한, 실제 DB에 가보면 테이블을 초기에 생성하지 않았지만 Python 코드에서 명시한 로직 때문에 테이블이 생성되어 있고 실제 데이터도 들어간 것을 확인할 수 있습니다.

이때, 입력된 content가 벡터로 변환되어 embed 컬럼에 vector 형태로 저장된 것도 확인할 수 있습니다.

---

FastAPI 개발: 데이터 조회하기

이제 벡터로 저장된 데이터 등을 조회할 수 있는 API를 구축하겠습니다. 

아래 API는 특정 title 값을 기반으로 텍스트 데이터를 검색하는 로직을 수행합니다. /search/title 엔드포인트에 get 방식으로 요청이 들어오면 로직을 수행하게 되는데요. text_item 테이블에서 title값이 요청 받은 text와 정확히 일치하는( Exact Match ) 데이터를 가지고 오게 됩니다.

@router.get("/search/title", response_model=List[TextItemResponse], tags=["Search"])
async def search_by_title(title: str, db: AsyncSession = Depends(get_db_session)):
    """
    title이 Exact Match인 레코드를 리스트로 반환
    """
    # 1) SELECT 쿼리 준비
    stmt = select(TextItem).where(TextItem.title == title)
    
    # 2) 실행
    result = await db.execute(stmt)
    
    # 3) Query 결과
    items = result.scalars().all()
    
    return items

 

그럼 위 코드의 실제 결과를 확인해볼까요? FastAPI에서 제공해주는 Swagger UI에서 실행을 해보도록 하겠습니다.

 

검색한 title을 '이수진입니다.'라는 제목으로 검색을 했고, 그 결과 3개의 결과 값이 나왔습니다. 3개의 데이터는 똑같이 '이수진입니다.'라는 제목을 가지고 있습니다. 다만, content 내용은 전부 다르죠. 이와 같이 위 API는 title이 정확히 일치 되는 것들을 가져오도록 수행하게 됩니다.

 

300x250

이제 마지막으로 content에 대한 벡터 서치(vector search)를 진행하는 코드를 살펴보겠습니다. 시맨틱 서치(semantic search)라고 할 수도 있는 이 로직은 사실 vector database로 사용하기 위한 핵심적인 로직이죠.

아래 API 코드는 입력 받은 쿼리(query)를 기반으로 의미적으로 유사한 텍스트를 코사인 유사도(cosine similarity)로 유사도 검색하여 결과를 제공하는 로직을 수행합니다.

@router.get("/search/semantic", tags=["Search"])
async def semantic_search(query: str, limit: int = 5, db: AsyncSession = Depends(get_db_session)):
    """
    1) query 문자열을 임베딩 (HuggingFace Hub)
    2) pgvector 메서드 방식 (cosine_distance)로 오름차순 정렬
    3) 상위 limit개 반환
    """
    # (1) 임베딩
    query_embedding = get_embedding(query)

    # (2) 검색
    stmt = (
        select(TextItem)
        .order_by(TextItem.embed.cosine_distance(query_embedding))
        .limit(limit)
    )
    # 쿼리 실행
    result = await db.execute(stmt)

    # (3) 결과 추출
    items = result.scalars().all()

    # (4) dict로 변환
    return [
        {
            "id": r.id,
            "title": r.title,
            "content": r.content
        }
        for r in items
    ]

 

사용자가 입력한 Query가 입력으로 들어오면 이를 앞서 정의한 임베딩 모델을 활용해 임베딩 벡터(embedding vector)로 변환해줍니다. 그리고 변환된 임베딩 벡터와 text_item 테이블이 가지고 있는 임베딩 벡터 컬럼인 embed와 코사인 유사도를 측정해 결과를 가져오게 되죠.

마찬가지로 실제 결과를 살펴보겠습니다. Swagger UI에서 실행시킨 결과는 다음과 같습니다.

 

 

저는 입력 쿼리로 '공부하고 있어요'라는 쿼리를 전달했고 limit은 1개 즉, 가장 유사도가 큰 1개만 가져오도록 했습니다.

그리고 그 결과 앞서 넣은 데이터 중 '열심히 공부하고 있습니다'와 의 결과가 나오게 되었죠.

 

이렇게 FastAPI를 구성하면 사용자의 입력에 따라 벡터 데이터베이스에 입력 데이터를 임베딩 모델을 활용해 벡터로 변환 후 저장할 수 있고 코사인 유사도와 같은 방법으로 벡터 검색(vector search)를 수행할 수 있습니다.

---

마무리

이번 포스팅은 pgvector 벡터 데이터베이스(vector database)를 활용해 파이썬 FastAPI와 연동하는 방법을 정리한 포스팅입니다.

허깅페이스의 embedding model을 활용해 사용자의 입력 데이터를 vector로 변환하여 저장하고, 벡터 서치와 타이틀 기반 완전 일치 검색을 수행할 수 있는 방법을 정리하였습니다.

긴 글이지만, 파이썬으로 벡터 데이터베이스를 연동하고 활용하는 방법이 궁금하신 분들에게 도움이 되길 바랍니다. 

감사합니다.