Mistral AI와 LlamaIndex로 ReAct 에이전트 만들기

MistralAI LLM을 사용해 도구(tool)와 RAG 파이프라인 위에서 ReAct 에이전트와 FunctionCalling 에이전트를 만드는 방법을 보여주는 노트북이에요. 간단한 계산기 도구부터 10K 문서 기반 RAG 도구까지, 에이전트에 도구를 붙이는 다양한 방식을 다룹니다.

출처: 문서

본문

Colab에서 이 노트북을 연다면 먼저 LlamaIndex를 설치해야 해요.

!pip install llama-index
!pip install llama-index-llms-mistralai
!pip install llama-index-embeddings-mistralai

import os
os.environ['MISTRAL_API_KEY'] = 'YOUR MISTRAL API KEY'

import json
from typing import Sequence, List

from llama_index.llms.mistralai import MistralAI
from llama_index.core.llms import ChatMessage
from llama_index.core.tools import BaseTool, FunctionTool
import nest_asyncio
nest_asyncio.apply()

에이전트용 계산기 도구 정의하기

에이전트에 붙일 간단한 계산기 도구를 만들어 봅니다.

def multiply(a: int, b: int) -> int:
    """Multiple two integers and returns the result integer"""
    return a * b

multiply_tool = FunctionTool.from_defaults(fn=multiply)

def add(a: int, b: int) -> int:
    """Add two integers and returns the result integer"""
    return a + b

add_tool = FunctionTool.from_defaults(fn=add)

MISTRAL_API_KEY가 설정되어 있는지 확인하세요. 아니면 api_key 파라미터를 명시적으로 지정하면 됩니다.

llm = MistralAI(model="mistral-large-latest")

FunctionCalling 에이전트 사용하기

계산기 함수로 간단한 FunctionCalling 에이전트를 초기화합니다.

from llama_index.core.agent import FunctionCallingAgentWorker
from llama_index.core.agent import AgentRunner

agent_worker = FunctionCallingAgentWorker.from_tools(
    [multiply_tool, add_tool],
    llm=llm,
    verbose=True,
    allow_parallel_tool_calls=False,
)
agent = AgentRunner(agent_worker)

response = agent.chat("What is (121 + 2) * 5?")
print(str(response))

# inspect sources
print(response.sources)

이번에는 병렬 함수 호출(parallel function calling)을 켜서 두 개의 곱셈 연산을 동시에 호출해 봅니다.

# enable parallel function calling
agent_worker = FunctionCallingAgentWorker.from_tools(
    [multiply_tool, add_tool],
    llm=llm,
    verbose=True,
    allow_parallel_tool_calls=True,
)
agent = AgentRunner(agent_worker)
response = await agent.achat("What is (121 * 3) + (5 * 8)?")
print(str(response))

ReAct 에이전트 사용하기

ReAct 에이전트도 동일한 도구로 만들 수 있어요.

from llama_index.core.agent import ReActAgent

agent = ReActAgent.from_tools([multiply_tool, add_tool], llm=llm, verbose=True)
response = agent.chat("What is (121 * 3) + (5 * 8)?")
print(str(response))

RAG 파이프라인 위의 에이전트

간단한 10K 문서 위에 Mistral FunctionCalling 에이전트를 구축해 봅니다. Mistral 임베딩과 mistral-medium으로 RAG 파이프라인을 만들고, 이를 Mistral 에이전트에 도구로 전달합니다.

!mkdir -p 'data/10k/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O 'data/10k/uber_2021.pdf'

from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.embeddings.mistralai import MistralAIEmbedding
from llama_index.llms.mistralai import MistralAI

embed_model = MistralAIEmbedding()
query_llm = MistralAI(model="mistral-medium")

# load data
uber_docs = SimpleDirectoryReader(
    input_files=["./data/10k/uber_2021.pdf"]
).load_data()

# build index
uber_index = VectorStoreIndex.from_documents(
    uber_docs, embed_model=embed_model
)
uber_engine = uber_index.as_query_engine(similarity_top_k=3, llm=query_llm)

query_engine_tool = QueryEngineTool(
    query_engine=uber_engine,
    metadata=ToolMetadata(
        name="uber_10k",
        description=(
            "Provides information about Uber financials for year 2021. "
            "Use a detailed plain text question as input to the tool."
        ),
    ),
)

from llama_index.core.agent import FunctionCallingAgentWorker
from llama_index.core.agent import AgentRunner

agent_worker = FunctionCallingAgentWorker.from_tools(
    [query_engine_tool], llm=llm, verbose=True
)
agent = AgentRunner(agent_worker)

response = agent.chat(
    "What are the risk factors for Uber in 2021?"
)
print(str(response))

같은 RAG 도구를 ReAct 에이전트로도 실행해 볼게요.

from llama_index.core.agent import ReActAgent

agent = ReActAgent.from_tools([query_engine_tool], llm=llm, verbose=True)
response = agent.chat("What are the risk factors for Uber in 2021?")
print(str(response))

더 알아보기 (Learn more)

LlamaIndex Agents 문서 — 에이전트 워크플로 가이드
FunctionCallingAgentWorker / AgentRunner — 함수 호출 기반 에이전트
ReActAgent — ReAct(Reason+Act) 패턴의 에이전트
mistral-large-latest — 이 예제에서 사용한 Mistral 모델