ColBERT와 OpenAI로 RAG 구현하기: LLM과 검색엔진을 결합한 대화형 정보 검색(IR) 시스템

이번 글에서는 ColBERT와 LLM을 이용해서 RAG를 구현해보려고 한다. 그리고 RAG 구현을 위해서 LLM 프롬프트를 어떻게 사용할 수 있는지 확인해 보자.

 

목표는 다음과 같다.

• RAG 구현을 위한 내부 모듈 및 시스템 파이프라인을 이해한다.
• ColBERT와 OpenAI API를 사용하여 RAG를 실제로 구현할 수 있다.
• Function calling을 포함한 프롬프트 엔지니어링 방법에 대해 이해한다.

이번 글에서 다룰 데이터셋은 LoTTE 벤치마크의 dev세트다. LoTTE 데이터 세트는 Wikipedia와 같은 엔티티 중심의 지식 기반에서는 잘 다루지 않을 수 있는 롱테일 주제에 대한 IR 시스템을 평가하기 위해 설계되었다. 이번 실습에서는 주로 반려 동물을 다루고 있는 문서를 사용한다.

 

1. 환경 설정

먼저 OpenAI 패키지를 설치한다. 그리고 RAG용 검색엔진으로 사용할 ColBERT 패키지를 설치한다. colbert 같은 경우, faiss-gpu를 사용하는 버전으로 설치해 줬다.

# OpenAI Python 패키지 설치
!pip install openai

# 검색엔진을 위해 ColBERT 사용
!pip install "colbert-ir[faiss-gpu, torch]"

설치가 완료되면 데이터셋을 다운받아주자.

from datasets import load_dataset

# LoTTE dataset 사용

dataset = 'lifestyle'
datasplit = 'dev'

collection_dataset = load_dataset("colbertv2/lotte_passages", dataset)
collection = [x['text'] for x in collection_dataset[datasplit + '_collection']]

queries_dataset = load_dataset("colbertv2/lotte", dataset)
queries = [x['query'] for x in queries_dataset['search_' + datasplit]]

print(f'Loaded {len(queries)} queries and {len(collection):,} passages')
print(queries[0])
print(collection[0])

2. RAG 개념 이해

RAG 개념을 이해하기 위해 먼저 레퍼런스 정보가 있다고 가정하고 그 정보를 토대로 질문에 답하는 예시를 실행해 보자.

# OpenAI 라이브러리 및 필요한 모듈을 가져온다.
import os
import json
from openai import OpenAI

# OpenAI API 키를 환경변수에서 설정한다.
os.environ["OPENAI_API_KEY"] = "Your API Key"

client = OpenAI()
# 사용할 모델을 설정해준다. 여기서는 gpt-3.5-turbo-1106 모델을 사용한다.
#llm_model = "gpt-4-1106-preview"
llm_model = "gpt-3.5-turbo-1106"

다음은 시뮬레이션을 위한 messages, reference, persona를 정의해 주자.

• messages는 임의로 식당 검색 도우미 어시를 만들어줬다.
• reference는 지금 당장 검색 API가 붙어있지 않기 때문에 가상으로 3개의 데이터를 만들어줬다.
• persona에는 Role을 식당 검색 도우미로 설정해 줬다.

messages = [
{"role": "user", "content": "너는 누구니"},
{"role": "assistant", "content": "저는 식당 검색 도우미 Good Place 라고 합니다."},
{"role": "user", "content": "강남역 근처 태국음식 추천해줘"}
]

reference = [
{"식당명": "파파야리프", "대표 메뉴": "똠양꿍", "분위기": "이국적/이색적", "기타": "데이트에 적합"},
{"식당명": "할랄가이즈", "대표 메뉴": "할랄", "분위기": "무슬림, 이국적", "기타": "가성비"},
{"식당명": "인더비엣", "대표 메뉴": "팟타이, 쌀국수", "분위기": "깔끔", "기타": "해장"}
]

persona = """
## Name: Good Place

## Role: 식당 검색 도우미

## Instruction
- 사용자의 이전 메시지 정보 및 주어진 검색 결과(JSON 형태로 제공) 정보를 활용하여 간결하게 답변을 생성한다.
- 주어진 검색 결과 정보로 대답할 수 없는 경우는 정보가 부족해서 답을 할 수 없다고 대답한다.
- 사용자가 사용한 언어로 답변을 생성한다.
"""

레퍼런스 없이 응답 생성

먼저 레퍼런스 없이 답변을 생성하는 경우를 살펴보자.

result = client.chat.completions.create(
    model=llm_model,
    messages=messages,
    temperature=0,
    seed=1
)

print(result.choices[0].message)

레퍼런스 없이 답변 생성

 

실제로 이렇게 응답은 하지만 '태국촌'이나 뒤에 나오는 음식점들은 실제로 존재하지 않는 음식점을 말해준다. 즉, hallucination이 발생했다.

 

레퍼런스 기반 응답 생성

그다음 레퍼런스를 토대로 답변을 생성하게 작성해 보자.

content = {
    "이전 메시지 정보": messages,
    "검색 결과": reference
}

msg = [{"role": "system", "content": persona}, {"role": "assistant",  "content": json.dumps(content, ensure_ascii=False, indent=4)}]

result = client.chat.completions.create(
    model=llm_model,
    messages=msg,
    temperature=0,
    seed=1
)

print(result.choices[0].message)

 

검색 결과(reference)를 활용하려면 LLM에 대화 히스토리와 검색 결과를 함께 전달해야 한다. 이렇게 하기 위해서 단순히 사용자와의 메시지를 전달하는 대신에 두 가지 데이터를 통합하여 하나의 JSON 형태로 만든다.

• messages: 사용자와 어시스턴트 간의 이전 대화 히스토리.
• reference: 검색 엔진에서 얻은 결과.

msg가 list형태로 값이 여러 개 들어가 있는 이유는 LLM에게 전달할 메시지의 구조와 역할이 각각 다르기 때문이다. LLM API는 메시지를 리스트 형태로 받아들이고, 각 요소는 다음 세 가지 역할 중 하나를 가질 수 있다:

• system: 시스템 메시지로, 모델의 역할과 동작 방식을 정의한다.
• user: 사용자 메시지로, 사용자가 모델에게 입력한 요청이나 질문을 나타낸다.
• assistant: 모델(LLM)이 생성한 답변이나 맥락 정보를 포함한다.

❓ 왜 현재 코드에는 user 항목이 없을까?

코드에서 user 메시지를 명시적으로 전달하지 않는 이유는:
• 사용자 메시지를 LLM이 단독으로 이해하지 않도록 하기 위함이다.
• 사용자의 요청과 검색 결과가 함께 전달되며, LLM은 이 정보를 맥락으로 이해하고 응답을 생성한다.

즉, 대화 맥락과 검색 결과가 별도로 전달되기보다는, 하나의 구조화된 데이터(content)로 통합되어 전달되고 있다.

쉽게 말하자면, 사용자 입력(예: {"role": "user", "content": "태국 음식 추천해 주세요."})을 별도로 전달하는 대신, 이 입력이 포함된 전체 대화 히스토리(messages)와 검색 결과(reference)를 묶어 하나의 데이터로 제공한다.

 

레퍼런스 기반 답변 생성

이렇게 답변이 왔다. 우리가 전에 넣어줬던 reference를 활용한 응답이라고 볼 수 있다.

 

이런 식으로 RAG는 레퍼런스가 주어졌을 때 LLM이 자체적으로 갖고 있는 지식 말고, 레퍼런스를 활용해 사실 기반 응답을 준다는 것을 알 수 있다.

 

3. 검색엔진 준비 - colBERT

그럼 이제 본격적으로 대화형 IR 쪽에서 검색엔진을 준비해 보자. 이번 실습에서는 colBERT를 활용해 볼 것이다.

빠른 실습을 위해서 위에서 다운받은 데이터셋 중에서 처음 10,000개의 구절에 대해서만 색인을 진행할 것이다.

# 필요한 라이브러리 import

from colbert import Indexer, Searcher
from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert.data import Queries, Collection
import colbert

# 테스트를 위해 전체 문서 중 10000개만 색인
checkpoint = 'colbert-ir/colbertv2.0'
index_name = "test"
doc_maxlen = 300
max_id = 10000

with Run().context(RunConfig(nranks=1, experiment='notebook')):  # nranks specifies the number of GPUs to use
    config = ColBERTConfig(doc_maxlen=doc_maxlen, nbits=2, kmeans_niters=4) # kmeans_niters specifies the number of iterations of k-means clustering; 4 is a good and fast default.                                                                           # Consider larger numbers for small datasets.

    indexer = Indexer(checkpoint=checkpoint, config=config)
    indexer.index(name=index_name, collection=collection[:max_id], overwrite=True)

위 작업은 대략 5분 정도 소요된다. 이렇게 색인이 완료되었다.

 

그럼 이제 검색을 위한 searcher를 생성해 보자.

with Run().context(RunConfig(experiment='notebook')):
    searcher = Searcher(index=index_name, collection=collection)

문서 검색

생성된 searcher를 통해 실제 검색을 진행해 보자. 데이터셋이 영어 문서이기 때문에 영어로 질의를 생성해줘야 한다.

영어 질의로 검색을 하고 5개를 가져오도록 설정했다.

queries = ["What are key factors to consider for minimizing stress when moving an aquarium?"]

for query in queries:
  print(f"#> {query}")

  # Find the top-5 passages for this query
  results = searcher.search(query, k=5)

  # Print out the top-k retrieved passages
  for passage_id, passage_rank, passage_score in zip(*results):
      print(f"\t [{passage_rank}] \t\t {passage_score:.1f} \t\t {searcher.collection[passage_id]}")

검색 결과

 

결과를 보면 답변을 잘 가져오는 걸 확인할 수 있다. 이렇게 1만 개의 데이터를 색인한 searcher를 준비해 봤다.

 

4. 대화형 IR

그럼 이제 준비된 검색엔진과 LLM을 활용해서 대화형 IR 시스템을 구현해 보자.

persona_qa = """
## Role: 반려 동물 전문가

## Instructions
- 사용자의 이전 메시지 정보 및 주어진 Reference(JSON 형태로 제공) 정보를 활용하여 간결하게 답변을 생성한다.
- 주어진 검색 결과 정보로 대답할 수 없는 경우는 다른 대답 하지 말고 "정보가 부족해서 답을 할 수 없습니다"라고 대답한다.
- 한국어로 답변을 생성한다.
"""

persona_function_calling = """
## Role: 반려 동물 전문가

## Instruction
- 사용자가 대화를 통해 반려 동물 관련된 지식이나 도움을 요청하면 search api를 호출할 수 있어야 한다.
- Search api에 필요한 standalone_query는 영어로 생성한다.
- 반려 동물과 관련되지 않은 나머지 대화 메시지에는 적절한 대답을 생성한다.
"""

tools = [
    {
        "type": "function",
        "function": {
            "name": "search",
            "description": "search knowledge from the user messages sent and received.",
            "parameters": {
                "properties": {
                    "standalone_query": {
                        "type": "string",
                        "description": "English query suitable for use in search from the user messages history."
                    }
                },
                "required": ["standalone_query"],
                "type": "object"
            }
        }
    },
]

 

4.1. 주요 페르소나

우선 두 개의 페르소나가 필요하다.

• persona_qa: 검색엔진의 결과를 기반으로 응답을 생성하는 역할.
  • 검색엔진으로부터 레퍼런스 정보가 뽑혔을 때 그 정보와 그전 메시지 히스토리를 기반으로 답을 만드는 페르소나다.
• persona_function_calling: 질문이 검색이 필요한 질문인지 판단하고 검색 기능을 유도하는 역할.
  • 첫 메시지가 들어왔을 때 그 메시지가 실제로 검색요청이 필요한 질문인지, 아니면 일반적인 대화로 진행되는 건지 판단한다. 만약에 검색엔진이 필요한 대화라면 search라는 function call이 유도되도록 하는 프롬프트다.

4.2 persona_function_calling과 tools의 상호작용

 

여기서 중요한 것은 persona_function_calling이 LLM에 행동 방침을 명시한다는 점이다. LLM은 이 지침을 기반으로 사용자의 메시지를 분석하고 standalone_query를 생성해야 한다는 작업 목표를 이해한다.

 

standalone_query 생성 과정

1. 사용자의 메시지를 보고 반려 동물 관련된 질문인지 판단한다.
2. 검색이 필요하다고 판단되면, 검색 엔진에 적합한 영어 쿼리를 생성한다.
3. 생성된 standalone_query는 tools에 정의된 search 함수로 전달된다.

4.3 tools의 역할

tools는 LLM이 외부 함수를 호출할 수 있도록 정의한 인터페이스다. 여기서는 search라는 이름의 함수를 정의해서 사용자가 입력한 질문에 대해 검색이 필요한 경우, 검색 쿼리를 생성하고 생성된 standalone_query를 외부 검색 엔진으로 전달할 수 있게 한다.

• standalone_query는 검색에 사용할 쿼리를 나타낸다. 이 쿼리는 영어로 작성되고 검색 엔진에서 사용할 수 있도록 설계되었다.

4.4 standalone_query 생성 흐름

standalone_query가 어떻게 생성되는지 구체적으로 설명하자면:

1. persona_function_calling의 역할:

• Instruction에 “검색 API에 필요한 standalone_query는 영어로 생성한다”라고 명시되어 있다.
• LLM은 이 지침을 기반으로 질문이 검색 요청인지 여부를 판단하고, 검색 요청이라면 검색 가능한 영어 문장으로 변환한다.

2. 사용자의 메시지를 분석:

• 예: "1살 영국 마스티프에게 얼마나 많은 양을 먹여야 하나요?"
• standalone_query로 변환: "How much food should I feed a 1-year-old English Mastiff?"

3. Function call 생성:

• standalone_query는 tools에 정의된 search 함수의 매개변수로 사용된다.

4.5 실제 구현

이제 conversation_search 함수에서 이를 어떻게 활용하는지 확인하자.

 

conversation_search 함수는 대화형 에이전트를 관리하는 역할을 한다. messages, persona, tools를 인자로 받는다. persona, tools는 위에서 정의한 것들이 들어오고 messaages는 실제로 주어진 걸 넘겨받는다.

import json

def conversation_search(messages, persona, tools):
    msg = [{"role": "system", "content": persona}] + messages
    result = client.chat.completions.create(
        model=llm_model,
        messages=msg,
        tools=tools,
        temperature=0,
        seed=1
    )
    response_message = result.choices[0].message

    if response_message.tool_calls:
        for tool_call in response_message.tool_calls:
            function_args = json.loads(tool_call.function.arguments)
            standalone_query = function_args.get("standalone_query")
            results = searcher.search(standalone_query, k=3)
            print("검색어>>: ", standalone_query)
            print()
            print("검색결과>>: ")
            retrieved_context = []
            for passage_id, passage_rank, passage_score in zip(*results):
                retrieved_context.append(searcher.collection[passage_id])
                print(f"\t [{passage_rank}] \t\t {passage_score:.1f} \t\t {searcher.collection[passage_id]}")

            msg = [{"role": "system", "content": persona_qa}] + messages
            msg.append({"role": "user", "content": json.dumps(retrieved_context)})
            qaresult = client.chat.completions.create(
                    model=llm_model,
                    messages=msg,
                    temperature=0,
                    seed=1
                )
            print("LLM 대답>>:")
            print(qaresult.choices[0].message.content)
    else:
        print("LLM 대답>>:")
        print(response_message.content)

주요 동작 흐름은 다음과 같다:

1. 메시지 초기화

conversation_search 함수에서 사용자 messages와 persona를 결합해서 LLM에 전달할 메시지를 생성한다. persona는 현재 에이전트의 역할과 행동 방침을 지정한다. 그리고 지금까지 메시지 히스토리를 이후 msg로 계속 추가해 준다.

2. LLM 호출

LLM 모델을 호출하여 사용자 메시지에 대한 초기 응답을 생성한다. 결과는 response_message에 저장한다.

3. 검색 기능 호출 판단

이후 결과는 두 가지로 나온다. 하나는 반려동물 관련 내용을 포함한 메시지 요청이면 tool_calls를 통해 function call이 유도되었다면 실제로 response_message 안에 tool_calls가 있게 된다. 그렇지 않으면 그냥 일반적인 대답을 유도했기 때문에 일반적인 LLM 대답이 출력될 것이다.

4. 검색 기능 실행

만약에 function call이 유도되었다면 결과를 json loads를 이용해 파싱 하고, standalone_query를 추출해 낸다. standalone_query를 통해서 아까 만든 searcher에 검색을 날리고 결과를 지금은 3개만 받아오도록 설정했다.

5. 검색 결과 기반 추가 응답 생성

3개 받아온 결과를 retrieved_context에 저장해 주고 아까 정의해 준 persona_qa를 통해서 한번 더 LLM을 통해서 답변을 생성하게 한다. 이때도 이전 메시지를 똑같이 추가해 준다. 마지막에 content를 한번 더 달아주게 된다. 그다음에 API Client를 호출하게 된다.

6. 응답 출력

최종 생성된 응답을 출력한다.

이제 실제 api를 실행해 보자. (아까 실행했지만 일단 한번 더 실행해 주자,,,)

# OpenAI 라이브러리 및 필요한 모듈을 가져옵니다.
import os
import json
from openai import OpenAI

# OpenAI API 키를 환경변수에서 설정합니다.
os.environ["OPENAI_API_KEY"] = "Your API Key"

client = OpenAI()
# 사용할 모델을 설정합니다. 여기서는 gpt-3.5-turbo-1106 모델을 사용합니다.
#llm_model = "gpt-4-1106-preview"
llm_model = "gpt-3.5-turbo-1106"

이제 그럴듯한 메시지를 입력해 보자.

강아지 관련 질문

messages = [
{"role": "user", "content": "안녕하세요, 저는 새로운 개 주인이고 방금 영국 마스티프를 입양했어요. 그를 돌보는 데 대해 몇 가지 팁을 줄 수 있나요?"},
{"role": "assistant", "content": "물론이죠! 영국 마스티프 같은 대형견을 돌보는 것은 매우 보람찬 일입니다. 그들은 균형 잡힌 식단, 규칙적인 운동, 그리고 정기적인 수의사 검진이 필요합니다. 그들의 돌봄에 대해 구체적으로 궁금한 점이 있나요?"},
{"role": "user", "content": "네, 특히 그의 식단에 대해 걱정이 되네요. 1살 영국 마스티프에게 얼마나 많은 양을 먹여야 하나요?"},
]

conversation_search(messages, persona_function_calling, tools)

검색 결과

 

일단 검색어를 잘 뽑아내서 search 하고 검색 결과를 3개 잘 가져온 걸 확인할 수 있다.

실제로 내용을 확인해 보면 한 살 영국 마스티프에게는 얼마나 많은 양을 먹여야 하는 거에 대한 정보는 잘 안 나온 걸 볼 수 있다. 그래서 LLM 대답은 정보가 부족해서 답을 할 수 없다고 나온다.

고양이 관련 질문

한번 더 다른 질문을 해보자.

messages = [
{"role": "user", "content": "저는 고양이를 입양하고 싶은데, 최근에 청각장애가 있는 고양이를 만났어요. 조금 걱정이 되네요. 특별히 고려해야 할 점이 있나요?"},
{"role": "assistant", "content": "청각장애가 있는 고양이를 입양하는 것은 독특하고 보람찬 경험이 될 수 있습니다. 이들은 일반적으로 안전한 실내 환경과 명확한 시각적 의사소통이 필요합니다. 특별히 걱정되는 점이 있나요?"},
{"role": "user", "content": "네, 그들의 행동에 대해 궁금해요. 청각장애가 있는 고양이는 다른 고양이들보다 더 공격적인가요?"}
]

conversation_search(messages, persona_function_calling, tools)

검색 결과

 

위처럼 검색어도 영어로 잘 뽑아냈고, 검색 결과도 3개 잘 가져왔다. 검색 결과에 실제 질문에 해당하는 내용이 있기 때문에 LLM 대답도 잘 답변하는 걸 볼 수 있다.

이렇게 reference를 활용해 잘 대답하는 걸 확인할 수 있다.

낙타 관련 질문

한번 더 다른 질문을 해보자.

messages = [
{"role": "user", "content": "낙타를 키우고 싶은데 어떤 준비가 필요한가요?"}
]

conversation_search(messages, persona_function_calling, tools)

검색 결과

 

검색 결과에 낙타에 관련된 내용이 없기 때문에 LLM 질문도 잘 못하는 걸 확인할 수 있다.

지식에 대한 질문이 아닌 경우

이번엔 지식에 대한 질문이 아닌 경우도 물어보자.

대화를 잘하고 있다가 이 LLM에 대해서 질문을 해보자.

messages = [
{"role": "user", "content": "저는 고양이를 입양하고 싶은데, 최근에 청각장애가 있는 고양이를 만났어요. 조금 걱정이 되네요. 특별히 고려해야 할 점이 있나요?"},
{"role": "assistant", "content": "청각장애가 있는 고양이를 입양하는 것은 독특하고 보람찬 경험이 될 수 있습니다. 이들은 일반적으로 안전한 실내 환경과 명확한 시각적 의사소통이 필요합니다. 특별히 걱정되는 점이 있나요?"},
{"role": "user", "content": "그런데 너는 누구니?"}
]

#conversation_search(messages, persona_function_calling, tools, {"type": "function", "function": {"name": "searcher"}})
conversation_search(messages, persona_function_calling, tools)

검색 결과

 

이전에는 검색어와 검색 결과를 기반으로 대답했다면, 이번에는 검색을 하지 않고 바로 LLM이 대답한 걸 확인할 수 있다.