티스토리 뷰

반응형

🚀 FastAPI 기반 로컬 LLM + RAG 챗봇 만들기 – 벡터 임베딩 저장 및 검색 구현

이번 글에서는 pgvector를 활용하여 FastAPI에서 벡터 데이터를 저장하고 유사도 검색을 수행하는 방법을 다룹니다.
LangChain을 사용한 벡터 임베딩 생성 → 벡터 데이터를 PostgreSQL에 저장 → 유사도 검색을 통한 RAG 구현 순서로 진행됩니다.


📌 1. PostgreSQL에 벡터 저장을 위한 테이블 생성

📌 PostgreSQL 내부에서 SQL 실행

CREATE TABLE IF NOT EXISTS documents (
    id SERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding vector(768) NOT NULL  -- 768차원 벡터 저장
);

설정 설명:

  • id SERIAL PRIMARY KEY → 자동 증가하는 기본 키
  • content TEXT NOT NULL → 원본 텍스트 저장
  • embedding vector(768) NOT NULL → 벡터(임베딩) 저장 (768차원)

📌 2. FastAPI에서 벡터 데이터를 PostgreSQL에 저장

반응형

📌 app/api/v1/endpoints/vector_store.py 파일을 생성합니다.

from fastapi import APIRouter, Depends
from pydantic import BaseModel
from sqlalchemy.orm import Session
from app.core.database import get_db
from app.core.embeddings import get_embedding
from sqlalchemy import text

router = APIRouter()

# 데이터 모델 정의
class DocumentRequest(BaseModel):
    content: str

@router.post("/store_vector/")
async def store_vector(request: DocumentRequest, db: Session = Depends(get_db)):
    # 텍스트를 벡터로 변환
    vector = get_embedding(request.content)

    # 벡터를 PostgreSQL에 저장
    sql = "INSERT INTO documents (content, embedding) VALUES (:content, :vector)"
    db.execute(text(sql), {"content": request.content, "vector": list(vector)})
    db.commit()

    return {"message": "Vector stored successfully"}

이제 FastAPI에서 /store_vector/ 엔드포인트를 JSON 형식으로 요청받아 처리할 수 있습니다! 🚀


📌 3. FastAPI에서 벡터 유사도 검색 구현

📌 app/api/v1/endpoints/vector_search.py 파일을 생성합니다.

from fastapi import APIRouter, Depends
from pydantic import BaseModel
from sqlalchemy.orm import Session
from app.core.database import get_db
from app.core.embeddings import get_embedding
from sqlalchemy import text

router = APIRouter()

# 데이터 모델 정의
class QueryRequest(BaseModel):
    query: str

@router.post("/search_vector/")
async def search_vector(request: QueryRequest, db: Session = Depends(get_db)):
    # 입력된 텍스트를 벡터로 변환
    query_vector = get_embedding(request.query)

    # 가장 유사한 벡터 찾기 (코사인 유사도 기준)
    sql = """
    SELECT content, embedding <=> CAST(:query_vector AS vector(768)) AS similarity
    FROM documents
    ORDER BY similarity ASC
    LIMIT 1;
    """
    result = db.execute(text(sql), {"query_vector": list(query_vector)}).fetchone()
    if result:
        return {"query": request.query, "most_similar_content": result[0], "similarity_score": result[1]}
    else:
        return {"message": "No matching results found"}

벡터 유사도 검색을 SQL 내장 함수(embedding <=> CAST(:query_vector AS vector(768)))를 활용하여 수행하도록 수정했습니다. 🚀


📌 4. FastAPI에 라우터 등록 (main.py 수정)

📌 FastAPI에 생성한 벡터 저장 및 검색 API를 등록합니다.

from fastapi import FastAPI
from app.api.v1.endpoints import vector_store, vector_search

app = FastAPI(title="FastAPI LLM Chatbot")

# 벡터 저장 및 검색 엔드포인트 추가
app.include_router(vector_store.router, prefix="/api/v1", tags=["vector_store"])
app.include_router(vector_search.router, prefix="/api/v1", tags=["vector_search"])

FastAPI 서버 실행 후 /api/v1/store_vector/ 및 /api/v1/search_vector/ API가 활성화됩니다! 🚀


📌 5. FastAPI 서버 실행 및 테스트

1️⃣ FastAPI 서버 실행

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

2️⃣ 벡터 데이터 저장 테스트 (curl 사용)

curl -X POST "http://localhost:8000/api/v1/store_vector/" \
     -H "Content-Type: application/json" \
     -d '{"content": "FastAPI는 Python 기반의 웹 프레임워크입니다."}'

📌 정상적으로 실행되었다면 데이터가 PostgreSQL에 저장됩니다!


3️⃣ 벡터 유사도 검색 테스트

curl -X POST "http://localhost:8000/api/v1/search_vector/" \
     -H "Content-Type: application/json" \
     -d '{"query": "FastAPI는 어떤 프레임워크인가요?"}'

📌 예상 응답:

{
    "query": "FastAPI는 어떤 프레임워크인가요?",
    "most_similar_content": "FastAPI는 Python 기반의 웹 프레임워크입니다.",
    "similarity_score": 0.01
}

유사도가 가장 높은 문서를 찾아 결과를 반환하는 것을 확인할 수 있습니다! 🚀


📌 6. 결론 및 다음 단계

pgvector를 활용하여 벡터 유사도 검색 기능 구현 완료!
FastAPI API 테스트 성공! 🚀

🚀 다음 글에서는 "3.1 Next.js 기반 프론트엔드 개발"을 시작합니다! 🚀

 

반응형
공지사항
최근에 올라온 글
최근에 달린 댓글
Total
Today
Yesterday
링크
«   2025/03   »
1
2 3 4 5 6 7 8
9 10 11 12 13 14 15
16 17 18 19 20 21 22
23 24 25 26 27 28 29
30 31
글 보관함
반응형