함수 호출로 여러 데이터베이스에 RAG 수행하기

함수 호출로 여러 데이터베이스에 RAG 수행하기 (RAG with multiple Databases via Function Calling)

함수 호출(function calling) 메커니즘으로 어떤 데이터베이스에 RAG 쿼리를 실행할지 결정하는 RAG 라우터 에이전트를 만드는 문서예요. RAG 로직 자체가 아니라 라우팅 부분을 다룹니다.

출처: 문서

본문

이 쿡북에서는 함수 호출 메커니즘을 사용해 어떤 데이터베이스에 RAG 쿼리를 실행할지 결정하는 RAG 에이전트를 코딩해요. 여기서는 RAG 로직 자체가 아니라 그 라우팅 부분만 구현해요. 이런 구조는 여러 가지 장점이 있어요:

여러 RAG 데이터베이스로 우아하게 확장할 수 있어요.
함수 호출의 "auto" 모드를 활성화해 RAG 단계와 클래식 LLM 지시 채팅 사이의 전환을 우아하게 처리해요.
이런 라우터는 평가(evaluation)와 파인튜닝이 아주 쉽습니다.

임포트와 Mistral 클라이언트

!pip install mistralai

from mistralai.client import Mistral
from getpass import getpass
import json

api_key= getpass("Type your API Key")
client = Mistral(api_key=api_key)

목 질문 (Mock Questions)

레모네이드 회사가 자체 내부 RAG를 만들고 싶다고 가정해요. 이 예시에서는 세 가지 유형의 데이터베이스에서 RAG 검색을 목(mock)으로 구현할게요:

HR -> 휴가, 혜택, 급여, 사무실 같은 인사 주제 정보를 담은 데이터베이스
Product -> 레모네이드 제품 특성(맛, 가격, 포장 색) 정보를 담은 데이터베이스
Finance -> 매출, 생산 비용, 부채 같은 재무 정보를 담은 데이터베이스

30개의 목 질문을 생성하는 함수를 만들어요.

def generate_questions():
    """
    Generate questions about HR, Prodcut, Finance or anything else. 
    
    Returns:
    List[str]: Llist of generated questions.
    """
    chat_response = client.chat.complete(
        model="mistral-large-latest",
        response_format={"type": "json_object"},
        temperature=1,
        messages=[
            {
                "role": "user",
                "content": """

### Context 
A Lemonade company wants to create its own internal RAG. 

This RAG is based on three data sources : 
 
- HR -> Database containing information about any human ressource topic e.g., holidays, pearks, salary, office
- Product -> Database containing information about the specificities of the Lemonade product e.g., available tastes, price, packaging color
- Finance -> Database containing financial information e.g. revenue, costs of production, liabilities 

### Task

Your role is too mock 30 questions related to either HR, Product, Finance or Other. These questions should reflect what new employee may ask. 

### Output format 

Your answer should take the form of a json contain a field "pairs" containg a list of lists of size 2 where the first element is a question and the second a label either "HR", "Product", "Finance" or "Other".
Questions of type "Other" should be totally random about things that have no link with the context.

Here is an example : 

{"pairs": [["How much money did the company made in 2024?","Finace"], ["How many days of holidays do I have ?","HR"]]}
"""
            }
        ]
    )
    return json.loads(chat_response.choices[0].message.content)

question_labels = generate_questions()

라우터 에이전트 구축 (Build router Agent)

함수 호출 도구 search_in_database를 정의하고, 질문에 따라 적절한 데이터 소스를 선택하는 시스템 프롬프트와 함께 라우터를 만들어요.

def get_response(question):
    """
    Generate questions about HR, Prodcut, Finance or anything else. 
    
    Returns:
    List[str]: Llist of generated questions.
    """

    tools = [
        {
            "type": "function",
            "function": {
                "name": "search_in_database",
                "description": "search_answer_in",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "question": {
                            "type": "string",
                            "description": "The question asked by used",
                        },
                        "source" : {
                            "type": "string",
                            "description": "Source to use to answer the question"
                        }
                    },
                    "required": ["source", "question"],
                },
            },
        }
    ]

    system_prompt = """
    Your are an AI assistant for a Lemonade company. 
    Your job is to help employee retrieve relevant information. 
    Specifically, you generate calls to a tool function which will then be able to perform efficient retrieval of relevant information.

    You dispose of a single tool function called "search_answer_in" which has two parameters "sources" and "question". 
    "question" is a copy paste of the user question. 
    "source" is a string whose values can exclusively be "HR", "Product", "Finance" or "Other" 

    The tool call will provide you some content to answer the user question. 
    Based on the the content from the tool call, answer the user's query. 

    You should be as helpful as possible while remaining factual, objective and keeping a professional tone. 
    You must answer in the same language as the user (or the one they ask you to).
    """

    chat_history = [
        {
            "role": "system",
            "content": system_prompt
        },
        {
            "role": "user",
            "content": question
        }
    ]

    chat_response = client.chat.complete(
        model="mistral-large-latest",
        temperature=0.3,
        messages=chat_history,
        tools=tools,
        tool_choice='any'
    )

    print(question)

    chat_history.append(chat_response.choices[0].message)

    db_to_search_in = json.loads(chat_response.choices[0].message.tool_calls[0].function.arguments)['source']

    print(f"Suggest search in {db_to_search_in} DB")
    
    return db_to_search_in

get_response('What are the different tastes of limonade ?')

라우터 에이전트 테스트 (Test router agent)

생성한 30개 질문에 대해 라우터가 올바른 데이터베이스를 선택하는지 테스트해요.

for question_label in question_labels['pairs'] : 
    print("##################") 
    searched_db = get_response(question_label[0])
    db_to_be_searched = question_label[1]
    if searched_db == db_to_be_searched:
        print("CORRECT")
    else:
        print("INCONRRECT") 
    print("")

질문이 모델에 전달되면 tool_choice='any'로 인해 항상 도구 호출이 발생하고, search_in_database 도구의 source 인자로 라우팅 대상 데이터베이스가 결정돼요. 이렇게 하면 RAG 검색이 실제로 어디에서 수행될지 미리 정확히 판별할 수 있어요.

함수 호출로 여러 데이터베이스에 RAG 수행하기