- 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}로 프롬프트 입력하여 응답 생성
전체 과정
-
-insert 옵션으로 문서 저장시, 파일 내용 읽고 Llama3로 임베딩 생성하여 PostgreSQL에 저장
-
-query 옵션 사용시, 프롬프트 임베딩 생성 후 items 테이블의 다른 임베딩들과 비교하여 "최근접" 문서 검색 (<-> 연산자로 L2 거리 계산)
- 검색된 문서를 프롬프트에 포함하여 Llama3에 전달, 채팅 응답 생성하여 출력
추가 팁
- 임베딩 생성에 특화된 모델 사용 고려 (llama3 대신)
- 영어 외 언어는 더 적합한 모델 탐색 필요
- L2 거리 외에 다른 거리 계산법 시도해볼 수 있음 (pgvector 가 다른 방법을 지원함)
- 전체 테이블 스캔은 확장성 낮으므로 pgvector 인덱스 등 활용
- 생성 단계에서 더 많은 문서 활용하거나, 키워드 매칭 등으로 추가 문서 가져오는 것도 도움됨
- 생성 프롬프트 조정 및 다양한 LLM 시도로 출력 품질 개선 가능