▲

PostgreSQL로 RAG 시스템 구축하기

 (pgdash.io)
26P by xguru 6달전 | favorite | 댓글과 토론
  • Retrieval-Augmented Generation(RAG) 시스템을 통해 기존의 내부 지식 기반(위키, 매뉴얼, 교육 및 참고 자료 등)을 바탕으로 질문에 답할 수 있는 AI 어시스턴트를 만들 수 있음
  • PostgreSQL, pgvector, ollama, 그리고 200줄 이하의 Go 코드만으로 RAG 시스템 구축 가능

Overview

  • 몇 단락의 이야기를 "문서 코퍼스"로 사용하고, 각 문서에 대해 Meta의 Llama3를 사용하여 문서 임베딩 생성(ollama로 로컬 호스팅)
  • 문서와 임베딩을 PostgreSQL 테이블에 저장하며, pgvector 확장을 통해 임베딩 저장 및 접근
  • 사용자 쿼리에 대해 가장 관련 있는 1개의 문서를 테이블에서 검색하여 Llama3를 사용해 응답 생성
  • ollama는 OpenAI와 유사한 HTTP API를 제공하여 임베딩과 채팅 응답을 생성해줌
  • Go 코드는 jackc/pgx와 pgvector-go를 사용하여 Postgres와 통신하며, ollama 클라이언트 API 패키지를 사용하여 HTTP API 호출을 처리함

Ollama로 모델들 실행하기

  • Ollama는 오픈소스 모델을 로컬에서 실행할 수 있게 해주는 도구로, OpenAI 스타일의 REST API를 제공함
  • ollama pull llama3 명령으로 llama3 모델 실행
  • ollama의 HTTP 서버는 기본적으로 127.0.0.1:11434에서 이용 가능

pgvector 설치

  • pgvector는 PostgreSQL 12~16 버전용 확장 기능으로, pgdg APT 저장소 사용시 sudo apt install postgresql-16-pgvector로 설치 가능
  • 설치 후 create extension vector;로 데이터베이스에 확장 기능 설정
  • create table items (id serial primary key, doc text, embedding vector(4096));로 문서와 임베딩 저장용 테이블 생성함

문서 데이터

  • Sherlock Holmes 이야기 "The Boscombe Valley Mystery"에서 4개의 단락을 사용함(퍼블릭 도메인-프로젝트 구텐베르크)

코드

  • GitHub에 MIT 라이선스로 공개된 데모 코드 사용 가능함
  • 문서 삽입은 INSERT INTO items (doc, embedding) VALUES ($1, $2)
  • 가장 관련성 높은 문서 검색은 SELECT doc FROM items ORDER BY embedding <-> $1 LIMIT 1 (<-> 연산자는 pgvector에 의해 제공)
  • Ollama API 호출은 ollama Go 패키지 사용함
    • 임베딩 생성은 api.EmbeddingRequest 사용
    • 채팅 응답 생성은 api.ChatRequest 사용 (프롬프트에 검색된 문서 포함)

커맨드라인 인터페이스

  • ragdemo -insert {path-to-doc-file}로 문서를 데이터베이스에 저장
  • ragdemo -query {query-text}로 프롬프트 입력하여 응답 생성

전체 과정

  1. -insert 옵션으로 문서 저장시, 파일 내용 읽고 Llama3로 임베딩 생성하여 PostgreSQL에 저장
  2. -query 옵션 사용시, 프롬프트 임베딩 생성 후 items 테이블의 다른 임베딩들과 비교하여 "최근접" 문서 검색 (<-> 연산자로 L2 거리 계산)
  3. 검색된 문서를 프롬프트에 포함하여 Llama3에 전달, 채팅 응답 생성하여 출력

추가 팁

  • 임베딩 생성에 특화된 모델 사용 고려 (llama3 대신)
  • 영어 외 언어는 더 적합한 모델 탐색 필요
  • L2 거리 외에 다른 거리 계산법 시도해볼 수 있음 (pgvector 가 다른 방법을 지원함)
  • 전체 테이블 스캔은 확장성 낮으므로 pgvector 인덱스 등 활용
  • 생성 단계에서 더 많은 문서 활용하거나, 키워드 매칭 등으로 추가 문서 가져오는 것도 도움됨
  • 생성 프롬프트 조정 및 다양한 LLM 시도로 출력 품질 개선 가능