저번 포스팅에서는 RAG의 개념과 탄생 배경에 대해 살펴보았다.
이번에는 실제로 Go와 Google Text Embedding API, SQLite를 이용해 검색 증강 생성(Retrieval-Augmented Generation)에서 Retrieval 부분을 구현해본다.
SQLite
SQLite는 파일 시스템이나 인메모리 데이터베이스로 사용할 수 있는 경량 데이터베이스이다.
이전 포스팅에서 설명했듯이 PostgreSQL, Mongo, Elastic Search 등을 사용해도 되지만, 다양한 구성을 직접적으로 피하기 위해 SQLite를 사용한다.
Google Text Embedding API
Google Text Embedding API는 Google에서 제공하는 텍스트 임베딩 API로, 문장을 벡터로 변환해준다.
Open AI나 다른 회사에서 제공하는 텍스트 임베딩 API도 활용할 수 있지만, 현재 Cowork AI에서 사용 중인 Google Text Embedding API를 사용한다.
최근 매우 성능이 향상된 새로운 모델인 text-embedding-large-exp-03-07을 발표하였고, 기존 차원 수 768에서 4배인 3072로 늘어나 성능이 크게 개선되었다고 하여 관심이 모아졌다.
해당 포스팅에서는 새로워진 모델의 결과도 확인해보도록 하자.
1. DB 셋업
SQLite를 Go에서 구동하기 위해서는 사전에 드라이버가 설치되어 있어야 한다.
go get github.com/mattn/go-sqlite3
그리고 sqlite_vec를 셋업하기 위해 다음 명령어를 실행한다.
go get -u github.com/asg017/sqlite-vec-go-bindings/cgo
이후 코드에서 다음과 같이 sqlite3을 구동한다. 첫 번째 파라미터는 드라이버 이름이며, 두 번째 파라미터는 읽어올 .db 파일의 경로이다.
파일 스토리지 기반으로 열 때는 특정 파일의 경로를 입력하면 되고, 인메모리 기반으로 열 때는 :memory:를 입력하면 된다.
db, err := sql.Open("sqlite3", "main.db")
if err != nil {
log.Panic(err)
}
이후 스키마를 작성한다. 필자는 구직 공고를 크롤링해왔기 때문에, job이라는 테이블을 만들었다.
-- in schema.sql
CREATE TABLE IF NOT EXISTS job(
id TEXT PRIMARY KEY,
company TEXT,
title TEXT,
content TEXT,
embedding BLOB
)
이후 go:embed 기능을 이용하여 SQL 파일을 읽어들이고 실행시킨다.
// in main.go
import _ "embed"
//go:embed schema.sql
var schema string
func initSchema(ctx context.Context, db *sql.DB) error {
if _, err := db.ExecContext(ctx, schema); err != nil {
return err
}
return nil
}
2. 검색 대상 문서를 가져오기
JobStore를 만들기 전에 init 메서드 내에 sqlite_vec.Auto()를 명시하였는데, 이는 추후 사용할 sqlite_vec 내장 함수를 활용하기 위함이다.init 함수는 패키지가 최초 로드될 때 1회만 호출된다.
크롤링된 문서는 이미 DB에 저장되어 있다고 가정한다.
QueryContext 메서드를 통해 JobStore와 생성, 조회를 담당하는 함수를 정의한다.
여기서 GetJobs는 단순히 모든 job을 가져온다. (물론 실제로는 필터링에 따라 적절한 쿼리가 필요할 수 있다.)
// in store/job.go
package store
import (
"context"
"database/sql"
sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/cgo"
"github.com/pkg/errors"
)
func init() {
sqlite_vec.Auto()
}
type Job struct {
ID string
Title string
Company string
Content string
Embedding []byte
Distance float64 // 보여주기 위한 코사인 거리 (저장 X)
}
type Jobs []Job
type JobStore interface {
GetJobs(ctx context.Context) (Jobs, error)
InsertOrReplaceJob(ctx context.Context, job Job) error
Search(ctx context.Context, query []byte) (Jobs, error)
}
type jobStore struct {
db *sql.DB
}
func NewJobStore(db *sql.DB) JobStore {
return &jobStore{db: db}
}
func (j jobStore) GetJobs(ctx context.Context) (Jobs, error) {
rows, err := j.db.QueryContext(ctx,
"SELECT id, title, company, content, embedding FROM job",
)
if err != nil {
return nil, errors.Wrap(err, "failed to get jobs")
}
defer rows.Close()
var jobs Jobs
for rows.Next() {
var job Job
if err = rows.Scan(&job.ID, &job.Title, &job.Company, &job.Content, &job.Embedding); err != nil {
return nil, errors.Wrap(err, "failed to scan job")
}
jobs = append(jobs, job)
}
return jobs, nil
}
3. Embedding API로 기존 Job 벡터화하기
이제 Embedding API를 통해 기존 Job을 벡터화한다.
Google Embedding API를 사용하기 위해서는 Google Cloud의 aiplatform을 사용해야 한다.
go get cloud.google.com/go/aiplatform
PredictionClient의 Predict() 메서드를 사용해 볼 텐데, 사용법은 다음과 같다.
3-1. Embedder 생성
예제 코드를 그대로 사용할 수도 있겠지만, 너무 길고 이해하기 어려울 수 있어 Embedder 코드를 따로 분리했다.
복잡해 보이지만 대부분의 코드는 구글 공식 레퍼런스의 예시를 참조하였다.
// in embed/embedder.go
package embed
import (
"context"
"fmt"
aiplatform "cloud.google.com/go/aiplatform/apiv1"
"cloud.google.com/go/aiplatform/apiv1/aiplatformpb"
"github.com/pkg/errors"
"google.golang.org/api/option"
"google.golang.org/protobuf/types/known/structpb"
)
type Embedder interface {
Embedding(ctx context.Context, taskType string, dimensionality int, input string) ([]float32, error)
}
type embedder struct {
client *aiplatform.PredictionClient
projectID string
location string
model string
}
func NewEmbedder(ctx context.Context, location string, projectID string, model string) (Embedder, error) {
apiEndpoint := fmt.Sprintf("%s-aiplatform.googleapis.com:443", location)
client, err := aiplatform.NewPredictionClient(ctx, option.WithEndpoint(apiEndpoint))
if err != nil {
return nil, errors.Wrap(err, "failed to create aiplatform client")
}
return &embedder{
client: client,
projectID: projectID,
location: location,
model: model,
}, nil
}
func (e embedder) Embedding(ctx context.Context, taskType string, dimensionality int, input string) ([]float32, error) {
endpoint := fmt.Sprintf("projects/%s/locations/%s/publishers/google/models/%s", e.projectID, e.location, e.model)
instances := []*structpb.Value{
structpb.NewStructValue(&structpb.Struct{
Fields: map[string]*structpb.Value{
"content": structpb.NewStringValue(input),
"task_type": structpb.NewStringValue(taskType),
},
}),
}
params := structpb.NewStructValue(&structpb.Struct{
Fields: map[string]*structpb.Value{
"outputDimensionality": structpb.NewNumberValue(float64(dimensionality)),
},
})
req := &aiplatformpb.PredictRequest{
Endpoint: endpoint,
Instances: instances,
Parameters: params,
}
resp, err := e.client.Predict(ctx, req)
if err != nil {
return nil, errors.Wrap(err, "failed to predict")
}
embeddings := make([][]float32, len(resp.Predictions))
for i, prediction := range resp.Predictions {
values := prediction.GetStructValue().Fields["embeddings"].GetStructValue().Fields["values"].GetListValue().Values
embeddings[i] = make([]float32, len(values))
for j, value := range values {
embeddings[i][j] = float32(value.GetNumberValue())
}
}
return embeddings[0], nil
}
3-2. Embedder를 사용하여 DB에 Embedding 저장하기
이제 Embedder를 사용해 Job을 벡터화해보자. 검색 대상은 Title과 Content로 하였다. 따라서 아래와 같이 코드를 작성할 수 있다.
여기서 저장 대상인 백터의 타입은 RETRIEVAL_DOCUMENT이다. 이는 정보 검색에 최적화된 임베딩을 생성하는데 특화되어 있다.
이외에도 타입 관련 옵션은 다양하게 준비되어 있으니, 공식 레퍼런스를 참고하여 사용에 맞는 유형을 선택하면 된다.
필자는 문서 벡터 저장 시 RETRIEVAL_DOCUMENT를 사용한다.
먼저 JobStore에 저장하는 함수를 추가한다.
SQLite는 INSERT OR REPLACE 문을 제공한다. 다른 DB에서는 UPSERT나 MERGE라고도 하며, 존재하지 않는 경우 INSERT, 존재하는 경우 UPDATE를 수행한다.
필자는 이를 이용해 생성과 업데이트를 동시에 처리하는 함수를 다음과 같이 작성했다.
func (j jobStore) InsertOrReplaceJob(ctx context.Context, job Job) error {
if _, err := j.db.ExecContext(ctx,
"INSERT OR REPLACE INTO job (id, title, company, content, embedding) VALUES (?, ?, ?, ?, ?)",
job.ID, job.Title, job.Company, job.Content, job.Embedding,
); err != nil {
return errors.Wrap(err, "failed to insert or replace job")
}
return nil
}
그리고 이제 임베딩을 얻은 후, Job의 임베딩 필드를 업데이트하는 코드를 작성한다.
// in main.go
// ...
for i, job := range jobs {
// 이미 임베딩이 있는 경우 스킵 (임베딩 모델 호출 비용이 크기 때문에)
if job.Embedding != nil {
slog.InfoContext(ctx, "skip job vector update",
"id", job.ID,
"progress", fmt.Sprintf("%d/%d", i+1, len(jobs)),
)
continue
}
// 임베딩 생성
embedding, err := embedder.Embedding(ctx, "RETRIEVAL_DOCUMENT", 3072, job.Title+" "+job.Content)
if err != nil {
log.Fatal(err)
}
// job 필드에 직렬화된 임베딩 업데이트
job.Embedding, err = sqlite_vec.SerializeFloat32(embedding)
if err != nil {
log.Fatal(errors.Wrap(err, "failed to serialize embedding"))
}
// job 업데이트
if err := jobStore.InsertOrReplaceJob(ctx, job); err != nil {
log.Fatal(err)
}
slog.InfoContext(ctx, "job vector updated",
"id", job.ID,
"progress", fmt.Sprintf("%d/%d", i+1, len(jobs)),
)
}
이렇게 하면 DB에 임베딩이 저장된 것을 확인할 수 있다.
4. query를 벡터화하기
이후 사용자의 질의, 즉 query를 벡터화해야 한다. 이를 위해 다음과 같이 코드를 작성할 수 있다.
테스트를 위해 질의를 사용자가 질문하는 형태로 작성하고, 타입은 QUESTION_ANSWERING으로 지정하여 산문 형태로 작성하였다.
// in main.go
// ...
query := "서버를 다루는 전문 개발자로서, RDB, MongoDB, ElasticSearch 등에 대한 경험을 보유하고 있습니다. 또한, AWS와 GCP와 같은 클라우드 서비스를 이용한 서버 구축 경험도 있습니다. 저에게 맞는 공고를 찾아주세요."
embedding, err := embedder.Embedding(ctx, "QUESTION_ANSWERING", 3072, query)
if err != nil {
log.Fatal("failed to embed query:", err)
}
queryVector, err := sqlite_vec.SerializeFloat32(embedding)
if err != nil {
log.Fatal(errors.Wrap(err, "failed to serialize query vector"))
}
5. 검색하기
이제 검색을 진행해보자. 여기서는 내장 함수 vec_distance_cosine를 사용해 코사인 거리를 계산한다.ASC로 정렬하면 관련성이 높은, 즉 가장 거리가 가까운 문서가 나오고, DESC로 정렬하면 관련성이 낮은, 즉 거리가 먼 문서가 나온다.
JobStore에 Search 메서드를 추가한다.
// in store/job.go
type JobStore interface {
// ...
Search(ctx context.Context, query []byte) (Jobs, error)
}
// ...
func (j jobStore) Search(ctx context.Context, query []byte) (Jobs, error) {
rows, err := j.db.QueryContext(ctx, `
SELECT
j.id,
j.title,
j.company,
j.content,
j.embedding,
vec_distance_cosine(j.embedding, ?) as distance
FROM job j
ORDER BY distance ASC
LIMIT 10;
`, query)
if err != nil {
return nil, errors.Wrap(err, "failed to search jobs")
}
defer rows.Close()
var jobs Jobs
for rows.Next() {
var job Job
if err = rows.Scan(&job.ID, &job.Title, &job.Company, &job.Content, &job.Embedding, &job.Distance); err != nil {
return nil, errors.Wrap(err, "failed to scan job")
}
jobs = append(jobs, job)
}
return jobs, nil
}
이후 다음과 같이 검색을 수행할 수 있다.
// in main.go
// ...
queryVector, err := sqlite_vec.SerializeFloat32(embedding)
if err != nil {
log.Fatal(err)
}
jobs, err = jobStore.Search(ctx, queryVector)
if err != nil {
log.Fatal(err)
}
for i, job := range jobs {
fmt.Println(fmt.Sprintf("%d.", i),
"ID:", job.ID,
"Company:", job.Company,
"Title:", job.Title,
"Length of Content:", len(job.Content),
"Distance:", job.Distance,
)
}
// ...
이렇게 하면 Retrieval-Augmented Generation의 R인 Retrieval을 구현할 수 있다.
결과
0. ID: d31fa2dd-5923-4381-84ac-520b92f40713 Company: 배달의민족 Title: [Tech] 로봇딜리버리플랫폼팀 서버 개발자 Length of Content: 3873 Distance: 0.23126929998397827
1. ID: 500d06b1-82b4-413e-b37c-29890243dccc Company: 뱅크샐러드 Title: 서버 엔지니어 (Server Engineer) Length of Content: 6501 Distance: 0.24647511541843414
2. ID: 6653e247-228c-420b-b902-761880d89978 Company: 네이버 Title: [NAVER] 광고 플랫폼 개발 (경력) Length of Content: 3673 Distance: 0.2595069408416748
3. ID: b2ab0026-b31c-476a-93b2-9ee03eddab2c Company: 배달의민족 Title: [Tech] CS 프로덕트실 백엔드 개발자 Length of Content: 4322 Distance: 0.26557084918022156
4. ID: 025a2ba9-7a35-47b7-8a62-a3e7b989598b Company: 카카오 Title: AI모델 플랫폼 개발 Length of Content: 3246 Distance: 0.2677247226238251
5. ID: 9defd036-6d00-48c1-9d43-18fb9f887a2e Company: 카카오 Title: DKOS(Kubernetes as a Service) 개발자 (경력) Length of Content: 3819 Distance: 0.2718163728713989
6. ID: fc5057e1-f658-4605-a327-fe0b56c09b26 Company: 네이버 Title: [NAVER] 컨테이너 기술 개발 (경력) Length of Content: 5755 Distance: 0.272807776927948
7. ID: 2d247818-ef6d-463e-801a-fb281a1255dc Company: 오늘의집 Title: [집중채용] Senior Software Engineer, Backend Length of Content: 6254 Distance: 0.2729299068450928
8. ID: fde0e8fd-0cbb-445c-b2c7-c9fc58f832fd Company: 카카오 Title: [공동체] 카카오게임즈 서비스 개발자(FE/BE) Length of Content: 2198 Distance: 0.2742129862308502
9. ID: af3e0e63-ded6-4fe9-bcf8-57bdce4735b9 Company: 두나무 Title: 클라우드 보안 담당자 Length of Content: 3066 Distance: 0.2753821313381195
이렇게 하면 생각보다 질의자가 알고 싶어 했던 서버 개발자에 적합한 공고들을 보여준다.
혹시 관련 데이터만 넣은 게 아닌지 의심할 수도 있지만, DESC로 정렬 기준을 변경하고 Distance도 추가로 로깅하면 다음과 같은 결과를 확인할 수 있다.
0. ID: 7cb7e73e-f91d-4242-a035-21730cc4eb90 Company: 라인 Title: Customer Care AI Project Associate Manager Length of Content: 1570 Distance: 0.38529109954833984
1. ID: 27eaab37-2048-4440-86ed-0430d1f83748 Company: 라인 Title: 일본어 전문 통번역사 (상시 채용) Length of Content: 2174 Distance: 0.36668115854263306
2. ID: dc46b690-7757-4eb7-beed-72a2dfce42ce Company: 라인 Title: Android Developer Length of Content: 1572 Distance: 0.3642077147960663
3. ID: 837e64ac-0772-49d5-ac11-a0b0d53df925 Company: 라인 Title: HR Operation (보훈대상자 특별채용) NEW Length of Content: 4220 Distance: 0.36187198758125305
4. ID: 6d44cb36-9fce-4276-8626-95931403c6c0 Company: 강남언니 Title: Talent Experience Manager Length of Content: 4997 Distance: 0.3570791482925415
5. ID: a1730346-2ea0-43d6-bb16-5980529ebb10 Company: 두나무 Title: 글로벌 마켓 리서치 인턴 Length of Content: 3607 Distance: 0.3563074767589569
6. ID: 2eabc4d1-21fc-4aff-a823-1bd2f7ce0a27 Company: 컬리 Title: [컬리] HMR MD Length of Content: 5009 Distance: 0.3562372922897339
7. ID: e1eec84a-c0a2-4d1f-890c-3384790383b7 Company: 당근 Title: Accounting Manager - 재무 Length of Content: 2506 Distance: 0.35504525899887085
8. ID: ab758bb0-cb33-4ee6-b4aa-7f054258c26e Company: 무신사 Title: Business Analyst (커머스 BA팀) Length of Content: 5931 Distance: 0.35138028860092163
9. ID: a23ed914-9946-42c6-a2dc-552efcaf150f Company: 야놀자 Title: Digital Communications Manager Length of Content: 6341 Distance: 0.3509758412837982
예상대로 서버 개발자와는 거리가 먼 일본어 통번역가 등의 포지션이 노출된다.
그렇다면 관련 없는 질문의 경우는 어떨까?
query := "하늘은 왜 파랗지?"
1. ID: 4c33e6b4-1d85-4df7-b48f-0c31cd8e88b2 Company: 야놀자 Title: MS PowerBI Engineer (1년, 육아휴직대체) Length of Content: 5337 Distance: 0.46892309188842773
2. ID: 1d1ebbb3-9c9e-41ba-a4fa-6255dda5adc7 Company: 네이버 Title: 2025 팀네이버 신입 공채 : Tech Length of Content: 97 Distance: 0.4691943824291229
3. ID: c3e1a0d0-239c-4504-9529-d2904d4701d6 Company: 야놀자 Title: [Yanolja Cloud Go Global] HR Manager Length of Content: 6413 Distance: 0.47171342372894287
4. ID: 193f8760-7bca-488f-a9a5-a0d2a0513880 Company: 야놀자 Title: 글로벌 인사기획 담당자 Length of Content: 6848 Distance: 0.4720369577407837
5. ID: cf53fcb0-9978-440d-9529-51e082c8b77e Company: 야놀자 Title: 인적 자원 관리 담당자 (1년, 계약직) Length of Content: 5033 Distance: 0.47699248790740967
6. ID: 6653e247-228c-420b-b902-761880d89978 Company: 네이버 Title: [NAVER] 광고 플랫폼 개발 (경력) Length of Content: 3673 Distance: 0.4846186637878418
7. ID: 1970e9fc-654d-47e8-a4eb-5b3fdec1ac98 Company: 강남언니 Title: 검색 프로덕트오너 Length of Content: 4091 Distance: 0.48492535948753357
8. ID: ef4376e5-ad58-4bb8-9503-9029b052f3bf Company: 오늘의집 Title: [집중채용] Senior Software Engineer, Frontend Length of Content: 6087 Distance: 0.48521751165390015
9. ID: af3e0e63-ded6-4fe9-bcf8-57bdce4735b9 Company: 두나무 Title: 클라우드 보안 담당자 Length of Content: 3066 Distance: 0.4853784739971161
그래도 결과는 나오며, TOP 10을 반환하므로 전혀 관련 없는 질문이라도 답이 나오게 된다.
다만, 주목할 점은 Distance 값의 차이이다. 서버 개발자 관련 질의에서는 가장 관련 있는 데이터가 약 0.2 정도의 거리를 보였으나, 관련 없는 질문에서는 가장 관련 있는 데이터조차 약 0.5에 가까운 거리가 산출된다.
실제 운영 시에는 여러 키워드로 사전 테스트를 진행한 후, 예를 들어 0.3 이상의 거리만 결과로 보이게 필터링하는 등의 방법을 사용할 수 있을 것이다.
마치며
단순한 검색 엔진을 넘어, 질의를 분석하여 적절한 벡터로 변환하는 것 자체가 불가능하다고 여겨졌던 시대가 있었음을 생각하면, 지금 이와 같은 구현은 매우 놀라운 발전이라 할 수 있다.
현재 많은 관심이 AI 쪽에 집중되고 있지만, 이와 같이 검색 및 데이터 처리 분야도 시간이 지남에 따라 끊임없이 발전하고 있다.
Reference
- https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings
- https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types
- https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings?hl=ko#generative-ai-get-text-embedding-go
- https://cloud.google.com/go/docs/reference/cloud.google.com/go/aiplatform/1.0.0/apiv1
- https://github.com/asg017/sqlite-vec
- https://alexgarcia.xyz/sqlite-vec/go.html
- https://www.sqlitetutorial.net/sqlite-replace-statement/