9. 개선: make_vector_store 명령에서 다수의 INSERT 쿼리를 묶어서 실행

9.1. bulk_create 적용

아래와 같이 개별 INSERT 쿼리로 실행하는 것보다

INSERT INTO document (page_content, metadata, embedding) VALUES ('...', '...', '...')
INSERT INTO document (page_content, metadata, embedding) VALUES ('...', '...', '...')
INSERT INTO document (page_content, metadata, embedding) VALUES ('...', '...', '...')

아래와 같이 묶어서 하나의 INSERT 쿼리로 실행하면, 데이터베이스와의 통신 횟수를 줄여 훨씬 빠르게 데이터를 저장할 수 있습니다.

INSERT INTO document (page_content, metadata, embedding)
       VALUES ('...', '...', '...'), ('...', '...', '...'), ('...', '...', '...')

장고 쿼리셋의 bulk_create(batch_size=None) 메서드를 활용하면 batch_size 개수만큼 하나의 INSERT 쿼리로 묶어줍니다. batch_size 인자를 생략하면 모든 문서를 한 번에 저장합니다. 한 번에 저장하는 개수가 너무 많으면 데이터베이스 메모리 사용량이 과도해져서 데이터베이스에서 오류를 발생할 수 있기에 레코드 개수가 많다면 batch_size 인자를 꼭 지정해주세요.

아래 코드는 기존의 개별 INSERT 쿼리로 실행되는 코드이구요.

for doc in tqdm(doc_list):
    paikdabang_menu_document = PaikdabangMenuDocument(
        page_content=doc.page_content,
        metadata=doc.metadata,
    )
    paikdabang_menu_document.save()

아래와 같이 모델 인스턴스를 리스트로 모아, 1000개씩 묶어서 데이터베이스로의 저장을 시도해봅니다.

# 객체만 생성할 뿐, 아직 데이터베이스 저장 전 입니다.
paikdabang_menu_documents = [
    PaikdabangMenuDocument(
        page_content=doc.page_content,
        metadata=doc.metadata,
    )
    for doc in doc_list
]

# 1000개씩 묶어서 데이터베이스로의 저장을 시도합니다.
PaikdabangMenuDocument.objects.bulk_create(paikdabang_menu_documents, batch_size=1000)

실행하면 아래와 같이 IntegrityError 예외가 발생합니다. embedding 컬럼은 NOT NULL 컬럼인데, 데이터베이스 저장 시에 embedding 컬럼에 값 지정없이 INSERT 쿼리가 수행되어 NULL 값으로 INSERT 쿼리가 수행되었구요. NOT NULL 제약 조건 위배로 예외가 발생했습니다. embedding 컬럼에 값이 지정되어 있었다면 예외없이 저장되었을 것입니다.

예외 발생: NOT NULL 제약 조건 위배

django.db.utils.IntegrityError: null value in column “embedding” of relation “chat_paikdabangmenudocument” violates not-null constraint DETAIL: Failing row contains (22, 1. 아이스티샷추가(아.샷.추) - SNS에서 더 유명…, {“source”: “빽다방.txt”}, null, 2025-01-31 09:10:36.083187+00, 2025-01-31 09:1207+00)

PaikdabangMenuDocument 모델에서는 django-lifecycle 훅을 통해 save 메서드 호출 전에 embedding 필드에 값을 지정하는 데요. 쿼리셋의 bulk_create 메서드는 각 인스턴스의 save() 메서드를 호출하지 않기 때문에, django-lifecycle hook을 통한 자동 임베딩 페이지에서 지정한 훅이 호출되지 않아 임베딩 값이 생성되지 않은 상황입니다.

9.2. bulk_create 시에 임베딩 값을 지원할려면?

쿼리셋의 bulk_create 메서드 호출 시에 임베딩 값이 지정되도록 할려면 어떻게 해야할까요? bulk_create 메서드를 재정의하여, 부모의 bulk_create 메서드를 호출하기 전에 .embedding 필드값을 지정하도록 해볼 수 있습니다. 첫번째 인자에는 앞서 생성했던 모델 인스턴스 리스트가 전달됩니다.

from typing import Iterable

class PaikdabangMenuDocumentQuerySet(models.QuerySet):
    def bulk_create(self, objs: Iterable["PaikdabangMenuDocument"], *args, **kwargs):
        # 각 모델 인스턴스마다 .embedding 필드에 임베딩 값 할당
        for obj in objs:
            obj.embedding = 계산된 임베딩 

        # 부모의 bulk_create 메서드 호출하여 데이터베이스에 저장
        return super().bulk_create(objs, *args, **kwargs)

objs 모델 인스턴스 리스트에서 각 모델 인스턴스마다 OpenAI 임베딩 API를 호출하는 것보다, 모아서 API 호출 횟수를 줄이면 네트워크 지연을 훨씬 줄일 수 있습니다. OpenAI 임베딩 API에서는 여러 문자열의 임베딩을 동시에 요청하는 기능도 제공해줍니다. 😜

  • str 타입의 값일 때에는 인자의 문자열 하나를 임베딩합니다. 각 임베딩 모델의 최대 토큰 수(예: 8191)를 초과해서는 안 됩니다.

  • List[str] 타입의 값으로 지정하여, 한 번의 요청으로 여러 텍스트의 임베딩을 동시에 요청할 수 있습니다. 리스트 내 각 문자열은 각 임베딩 모델의 최대 토큰 수(예: 8191)를 초과해서는 안 되며, 리스트 전체는 모델의 요청 제한(Rate Limit)을 초과하지 않는 범위에서 지원됩니다. Tier 1 계정일 경우 분당 최대 100만 토큰의 요청을 지원합니다. 그럼 8090 토큰을 가지는 문자열을 한 번에 최대 124개까지 요청할 수 있습니다.

OpenAI 각 모델의 요청 제한 수는 공식문서 Rate limits를 통해 확인하실 수 있습니다. 모델 별, 각 계정의 tier 별로 제한 수가 다릅니다.

text-embedding-3-small, text-embedding-3-large 모델의 요청 제한 수 (2025년 2월 기준)

Tier

RPM (분당 API 최대 요청수)

RPD (하루당 API 최대 요청수)

TPM (분당 최대 토큰수)

Batch Queue Limit

Free

100

2,000

40,000

-

Tier 1

3,000

-

1,000,000

3백만 토큰 (TPM*3배)

Tier 2

5,000

-

1,000,000

2천만 토큰 (TPM*20배)

Tier 3

5,000

-

5,000,000

1억 토큰 (TPM*20배)

Tier 4

10,000

-

5,000,000

5억 토큰 (TPM*100배)

Tier 5

10,000

-

10,000,000

40억 토큰 (TPM*400배)

Batch Queue Limit

Batch Queue Limit는 배치 요청 큐에 대기시킬 수 있는 최대 토큰 수입니다. Batch를 활용하면 실시간 임베딩 요청에 비해서 비용이 50% 절감되고, TPM 대비 3배~400배의 토큰 수를 한 번에 대기시킬 수 있습니다.

Batch에 대기시킬려는 토큰 수가 Batch Queue Limit을 초과한 Batch 요청은 아래 오류가 발생합니다.

Enqueued token limit reached for text-embedding-3-small in organization org-???. Limit: 20,000,000 enqueued tokens. Please try again once some in_progress batches have been completed.

OpenAI API 사용량이 많아지면, OpenAI 측에서 Tier를 한 단계씩 올려줍니다.

9.3. embed 함수에 리스트 지원 추가하기

PaikdabangMenuDocument 모델의 두 embed 함수에 리스트 지원을 추가합니다. OpenAI 임베딩 API 응답에서 response.data는 항상 리스트입니다.

  • input 인자로 문자열을 지정하면, 하나의 임베딩을 수행하고 response.data 는 벡터값 하나를 가지는 리스트를 반환합니다.

  • input 인자로 문자열 리스트를 지정하면, 여러 임베딩을 수행하고 response.data 는 다수의 벡터값을 가지는 리스트를 반환합니다.

인자로 문자열을 받으면 벡터값 하나를 반환하고, 문자열 리스트를 받으면 벡터값 리스트를 반환토록 변경하겠습니다.

 1from typing import List, Union
 2
 3class PaikdabangMenuDocument(LifecycleModelMixin, models.Model):
 4    # ...
 5
 6    @classmethod
 7    def embed(cls, input: Union[str, List[str]]) -> Union[List[float], List[List[float]]]:
 8        client = openai.Client(api_key=cls.openai_api_key, base_url=cls.openai_base_url)
 9        response = client.embeddings.create(
10            input=input,
11            model=cls.embedding_model,
12        )
13        if isinstance(input, str):
14            return response.data[0].embedding
15        return [v.embedding for v in response.data]
16
17    @classmethod
18    async def aembed(cls, input: Union[str, List[str]]) -> Union[List[float], List[List[float]]]:
19        client = openai.AsyncClient(api_key=cls.openai_api_key, base_url=cls.openai_base_url)
20        response = await client.embeddings.create(
21            input=input,
22            model=cls.embedding_model,
23        )
24        if isinstance(input, str):
25            return response.data[0].embedding
26        return [v.embedding for v in response.data]

9.4. bulk_create 메서드에 적용하기

이제 아래와 같이 objs 리스트에서 문자열 리스트를 생성한 후에, 벡터값을 생성/저장하고, 부모의 bulk_create 메서드를 호출하여 데이터베이스에 저장할 수 있습니다.

 1from typing import Iterable, List
 2
 3class PaikdabangMenuDocumentQuerySet(models.QuerySet):
 4    def bulk_create(self, objs: Iterable["PaikdabangMenuDocument"], *args, **kwargs):
 5        # 문자열 리스트 생성
 6        input_list: List[str] = [obj.page_content for obj in objs]
 7
 8        # 문자열 리스트를 벡터 리스트로 **한 번의 API 요청**으로 변환
 9        embedding_list: List[List[float]] = self.model.embed(input_list)
10
11        # 각 순서대로 개별 인스턴스에 벡터 값 할당
12        for obj, embedding in zip(objs, embedding_list):
13            obj.embedding = embedding
14
15        # 부모의 bulk_create 메서드 호출하여 데이터베이스에 저장
16        return super().bulk_create(objs, *args, **kwargs)

위 코드는 Rate Limit을 초과하지 않는 범위 내에서는 잘 동작합니다. 하나의 문자열에 대한 임베딩 토큰 수가 8090 일때, 124개 문자열을 임베딩 요청하면 총 토큰 수는 100만이 넘게 됩니다. Tier 2 계정일 경우 TPM(분당 최대 토큰수)이 100만 이므로, TPM 제한에 걸려 아래와 같은 RateLimitError 예외가 발생합니다.

예외 발생

RateLimitError: Error code: 429 - {‘error’: {‘message’: ‘Request too large for text-embedding-3-small in organization org-******************** on tokens per min (TPM): Limit 1000000, Requested 1003160. The input or output tokens must be reduced in order to run successfully. Visit https://platform.openai.com/account/rate-limits to learn more.’, ‘type’: ‘tokens’, ‘param’: None, ‘code’: ‘rate_limit_exceeded’}}

9.5. TPM 허용 범위 만큼 묶어서 임베딩 요청하기

Tier 1 계정일 경우 text-embedding-3-small 모델 TPM(분당 최대 토큰수) 제한이 1,000,000 이므로, 계정당 1분에 최대 1,000,000 토큰까지 임베딩할 수 있습니다. 각 계정의 TPM 제한은 공식문서를 통해서만 알 수 있을 뿐 API를 통한 조회는 지원하지 않기에, RAG_EMBEDDING_MAX_TOKENS_LIMIT 설정을 통해 직접 제한 설정을 두고 이 설정 값에 맞춰 그룹을 만들어 그룹 단위로 임베딩 요청하도록 하겠습니다.

mysite/settings.py
OPENAI_API_KEY = env.str("OPENAI_API_KEY", default=None)
RAG_EMBEDDING_MODEL = env.str("RAG_EMBEDDING_MODEL", default="text-embedding-3-small")
RAG_EMBEDDING_DIMENSIONS = env.int("RAG_EMBEDDING_DIMENSIONS", default=1536)
# Tier1, text-embedding-3-small 모델의 TPM : 1,000,000
RAG_EMBEDDING_MAX_TOKENS_LIMIT = env.int("RAG_EMBEDDING_MAX_TOKENS_LIMIT", default=1_000_000/10)

PaikdabangMenuDocument 모델에도 embedding_max_tokens_limit 클래스 변수를 추가하고, 디폴트 값으로 RAG_EMBEDDING_MAX_TOKENS_LIMIT 설정을 지정합니다.

chat/models.py
1class PaikdabangMenuDocument(LifecycleModelMixin, models.Model):
2    openai_api_key = settings.OPENAI_API_KEY
3    embedding_model = settings.RAG_EMBEDDING_MODEL
4    embedding_dimensions = settings.RAG_EMBEDDING_DIMENSIONS
5    embedding_max_tokens_limit = settings.RAG_EMBEDDING_MAX_TOKENS_LIMIT
6    # ...

임베딩 API에서는 문자열을 토큰으로 먼저 변환한 뒤에 임베딩 벡터로 최종 변환합니다. "hello, world" 문자열은 12글자이지만, text-embedding-3-small 모델에서 토큰은 [15339, 11, 1917]로서 3개가 되고, 임베딩 벡터는 1536차원으로서 [-0.01657603681087494, -0.03527357801795006, ...]로 생성됩니다.

토큰 수를 기반으로 여러 문자열들을 그룹으로 묶을려면, 각 문자열들을 토큰으로 변환하고 토큰 수를 계산하는 과정이 필요합니다. OpenAI에서는 토큰 수를 계산해주는 API는 제공하지 않습니다. OpenAI 공식문서 How to count tokens with Tiktoken에 따르면 tiktoken 라이브러리를 통해 API 호출없이도 토큰을 생성할 수 있다고 합니다.

PaikdabangMenuDocument 모델에 클래스 함수 get_token_size 메서드를 추가하여, 모델에 지정된 임베딩 모델을 기준으로 주어진 텍스트의 토큰 수를 계산하여 반환토록 하구요. PaikdabangMenuDocumentQuerySet에서 토큰 수 계산 시에 활용하겠습니다.

 1import tiktoken
 2
 3class PaikdabangMenuDocument(LifecycleModelMixin, models.Model):
 4    embedding_model = settings.RAG_EMBEDDING_MODEL
 5
 6    # ...
 7
 8    @classmethod
 9    def get_token_size(cls, text: str) -> int:
10        encoding: tiktoken.Encoding = tiktoken.encoding_for_model(cls.embedding_model)
11        token: List[int] = encoding.encode(text or "")
12        return len(token)

문자열 리스트를 인자로 받으면, 토큰 수에 기반하여 문자열 그룹을 생성해주는 make_groups_by_length 함수를 chat/utils.py 파일에 구현합니다.

쿼리셋의 bulk_create 메서드에서는 make_groups_by_length 함수를 활용하여 토큰 수 제한에 맞춰 문자열 리스트를 그룹핑하고, 각 그룹 별로 임베딩 API를 호출하여 임베딩 벡터를 생성합니다. 임베딩 API 호출 시에 Rate Limit 예외가 발생하면 60초 쉰 후에 최대 3번까지 재시도합니다.

 1import logging
 2import time
 3
 4from chat.utils import make_groups_by_length
 5
 6logger = logging.getLogger(__name__)
 7
 8class PaikdabangMenuDocumentQuerySet(models.QuerySet):
 9    # ...
10
11    def bulk_create(self, objs, *args, max_retry=3, interval=60, **kwargs):
12        # 임베딩된 벡터 데이터를 저장할 리스트
13        embeddings = []
14
15        groups = make_groups_by_length(
16            # 임베딩을 할 문자열 리스트
17            text_list=[obj.page_content for obj in objs],
18            # 그룹의 최대 허용 크기 지정
19            group_max_length=self.model.embedding_max_tokens,
20            # 토큰 수 계산 함수
21            length_func=self.model.get_token_size,
22        )
23
24        # 토큰 수 제한에 맞춰 묶어서 임베딩 요청
25        for group in groups:
26            for retry in range(1, max_retry + 1):
27                try:
28                    embeddings.extend(self.model.embed(group))
29                    break
30                except openai.RateLimitError as e:
31                    if retry == max_retry:
32                        raise e
33                    else:
34                        msg = "Rate limit exceeded. Retry after %s seconds... : %s"
35                        logger.warning(msg, interval, e)
36                        time.sleep(interval)
37
38        for obj, embedding in zip(objs, embeddings):
39            obj.embedding = embedding
40
41        return super().bulk_create(objs, *args, **kwargs)
42
43    # TODO: 비동기 버전 지원
44    async def abulk_create(self, objs, *args, max_retry=3, interval=60, **kwargs):
45        raise NotImplementedError
46        return await super().abulk_create(objs, *args, **kwargs)

문자열 리스트에서 토큰 수를 기반으로 그룹을 만들어주는 함수 make_groups_by_length를 아래와 같이 구현합니다.

chat/utils.py
 1from logging import getLogger
 2from typing import Callable, Generator, Iterable, List
 3
 4logger = getLogger(__name__)
 5
 6def make_groups_by_length(
 7    text_list: Iterable[str],
 8    group_max_length: int,
 9    length_func: Callable[[str], int] = len,
10) -> Generator[List[str], None, None]:
11    batch, group_length = [], 0
12    for text in text_list:
13        text_length = length_func(text)
14        if group_length + text_length >= group_max_length:
15            msg = "Made group : length=%d, item size=%d"
16            logger.debug(msg, group_length, len(batch))
17            yield batch  # 현재 배치 반환
18            batch, group_length = [], 0
19        batch.append(text)
20        group_length += text_length
21    if batch:
22        msg = "Made group : length=%d, item size=%d"
23        logger.debug(msg, group_length, len(batch))
24        yield batch  # 마지막 배치 반환

make_vector_store 명령을 수행해보시면, 빽다방.txt 파일에 대해서는 하나의 그룹만 생성이 되었구요. 이는 한 번의 임베딩 API 요청 만으로 임베딩을 수행했음을 의미합니다.

$ uv run python manage.py make_vector_store ./chat/assets/빽다방.txt
loaded 1 documents
split into 10 documents
100%|████████████████████████████████| 10/10 [00:00<00:00, 12409.18it/s]
[2025-02-02 10:41:22,525] Made group : length=854, item size=10