Arxiv에서 관련 연구 논문을 찾는 LLM 에이전트 만들기

Arxiv에서 관련 연구 논문을 찾는 LLM 에이전트 만들기

MistralAI 언어 모델을 기반으로 사용자 질의와 관련된 Arxiv 연구 논문을 찾고 요약하는 LLM 에이전트를 만드는 튜토리얼이에요. LlamaIndex 프레임워크로 구축하며, 에이전트는 RAG 쿼리 엔진·논문 가져오기·PDF 다운로드라는 세 가지 도구를 사용합니다.

출처: 문서

본문

에이전트가 사용하는 도구는 다음과 같아요.

  • RAG Query Engine: 최근 Arxiv 논문을 저장·검색해 지식 베이스로 활용해 관련 정보에 효율적이고 빠르게 접근합니다.
  • Paper Fetch Tool: RAG 쿼리 엔진에 없는 주제를 사용자가 지정하면 해당 주제의 최근 논문을 Arxiv에서 직접 가져옵니다.
  • PDF Download Tool: Arxiv가 제공하는 링크로 연구 논문의 PDF를 로컬에 다운로드하게 해 줍니다.

이 노트북은 Andrei Chernov가 만들었습니다. (Github, Linkedin)

먼저 필요한 라이브러리를 설치합니다. Mistral 모델에 접근하려면 API 키도 필요해요.

!pip install arxiv==2.1.3 llama_index==0.12.3 llama-index-llms-mistralai==0.3.0 llama-index-embeddings-mistralai==0.3.0

from getpass import getpass
import requests
import sys
import arxiv

from llama_index.llms.mistralai import MistralAI
from llama_index.embeddings.mistralai import MistralAIEmbedding
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Document, StorageContext, load_index_from_storage, PromptTemplate, Settings
from llama_index.core.tools import FunctionTool, QueryEngineTool
from llama_index.core.agent import ReActAgent

API 키는 여기에서 얻을 수 있어요.

api_key= getpass("Type your API Key")

llm = MistralAI(api_key=api_key, model='mistral-large-latest')

RAG 쿼리 엔진을 만들려면 임베딩 모델이 필요합니다. 이 튜토리얼에서는 MistralAI 임베딩 모델을 사용할게요.

model_name = "mistral-embed"
embed_model = MistralAIEmbedding(model_name=model_name, api_key=api_key)

이 튜토리얼을 무료 Mistral API 버전에서 접근 가능하게 유지하기 위해 최근 논문 10개만 내려받습니다. 더 많이 내려받으면 나중에 RAG 쿼리 엔진을 만들 때 한도를 초과할 수 있어요. Mistral 구독이 있다면 추가 논문을 받을 수 있습니다.

def fetch_arxiv_papers(title :str, papers_count: int):
    search_query = f'all:"{title}"'

    search = arxiv.Search(
        query=search_query,
        max_results=papers_count,
        sort_by=arxiv.SortCriterion.SubmittedDate,
        sort_order=arxiv.SortOrder.Descending
    )

    papers = []
    # Use the Client for searching
    client = arxiv.Client()
    # Execute the search
    search = client.results(search)
    for result in search:
        paper_info = {
            'title': result.title,
            'authors': [author.name for author in result.authors],
            'summary': result.summary,
            'published': result.published,
            'journal_ref': result.journal_ref,
            'doi': result.doi,
            'primary_category': result.primary_category,
            'categories': result.categories,
            'pdf_url': result.pdf_url,
            'arxiv_url': result.entry_id
        }
        papers.append(paper_info)
    return papers

papers = fetch_arxiv_papers("Language Models", 10)

[[p['title']] for p in papers]

이 과정은 임베딩 모델로 문서의 각 청크에 대한 벡터 표현을 만듭니다. Arxiv 포맷을 LlamaIndex가 이해하는 문서로 변환합니다.

def create_documents_from_papers(papers):
    documents = []
    for paper in papers:
        content = f"Title: {paper['title']}\n" \
                  f"Authors: {', '.join(paper['authors'])}\n" \
                  f"Summary: {paper['summary']}\n" \
                  f"Published: {paper['published']}\n" \
                  f"Journal Reference: {paper['journal_ref']}\n" \
                  f"DOI: {paper['doi']}\n" \
                  f"Primary Category: {paper['primary_category']}\n" \
                  f"Categories: {', '.join(paper['categories'])}\n" \
                  f"PDF URL: {paper['pdf_url']}\n" \
                  f"arXiv URL: {paper['arxiv_url']}\n"
        documents.append(Document(text=content))
    return documents

#Create documents for LlamaIndex
documents = create_documents_from_papers(papers)

Settings.chunk_size = 1024
Settings.chunk_overlap = 50

index = VectorStoreIndex.from_documents(documents, embed_model=embed_model)

인덱스 저장하기

많은 텍스트를 인덱싱하는 것은 임베딩 모델에 API 호출을 하므로 시간과 비용이 들 수 있어요. 실제 애플리케이션에서는 재인덱싱을 피하려고 벡터 데이터베이스에 인덱스를 저장하는 게 좋지만, 이 튜토리얼에서는 단순하게 벡터 데이터베이스 없이 인덱스를 로컬 디렉터리에 저장합니다.

index.storage_context.persist('index/')

# rebuild storage context
storage_context = StorageContext.from_defaults(persist_dir='index/')

#load index
index = load_index_from_storage(storage_context, embed_model=embed_model)

각 도구에 의미 있는 이름과 명확한 설명을 제공하는 것이 좋은 습관이에요. 에이전트가 필요할 때 가장 적절한 도구를 고르는 데 도움이 됩니다.

query_engine = index.as_query_engine(llm=llm, similarity_top_k=5)

rag_tool = QueryEngineTool.from_defaults(
    query_engine,
    name="research_paper_query_engine_tool",
    description="A RAG engine with recent research papers.",
)

RAG 도구가 문맥을 기반으로 질의에 답할 때 쓰는 프롬프트를 살펴봅니다. 참고로 두 개의 프롬프트가 있어요. 기본적으로 LlamaIndex는 답을 반환하기 전에 refine 프롬프트를 사용합니다. 응답 모드에 대한 더 자세한 내용은 여기를 참고하세요.

from llama_index.core import PromptTemplate
from IPython.display import Markdown, display

# define prompt viewing function
def display_prompt_dict(prompts_dict):
    for k, p in prompts_dict.items():
        text_md = f"**Prompt Key**: {k}" f"**Text:** "
        display(Markdown(text_md))
        print(p.get_template())
        display(Markdown(""))

prompts_dict = query_engine.get_prompts()
display_prompt_dict(prompts_dict)

나머지 두 도구는 단순히 Python 함수이므로 만들기 쉽습니다.

def download_pdf(pdf_url, output_file):
    """
    Downloads a PDF file from the given URL and saves it to the specified file.

    Args:
        pdf_url (str): The URL of the PDF file to download.
        output_file (str): The path and name of the file to save the PDF to.

    Returns:
        str: A message indicating success or the nature of an error.
    """
    try:
        # Send a GET request to the PDF URL
        response = requests.get(pdf_url)
        response.raise_for_status()  # Raise an error for HTTP issues

        # Write the content of the PDF to the output file
        with open(output_file, "wb") as file:
            file.write(response.content)
        return f"PDF downloaded successfully and saved as '{output_file}'."
    except requests.exceptions.RequestException as e:
        return f"An error occurred: {e}"

download_pdf_tool = FunctionTool.from_defaults(
    download_pdf,
    name='download_pdf_file_tool',
    description='python function, which downloads a pdf file by link'
)

fetch_arxiv_tool = FunctionTool.from_defaults(
    fetch_arxiv_papers,
    name='fetch_from_arxiv',
    description='download the {max_results} recent papers regarding the topic {title} from arxiv'
)

에이전트와 대화하기

세 도구로 ReAct 에이전트를 만듭니다.

# building an ReAct Agent with the three tools.
agent = ReActAgent.from_tools([download_pdf_tool, rag_tool, fetch_arxiv_tool], llm=llm, verbose=True)

ReAct 에이전트는 두 가지 주요 단계로 동작해요.

  • 추론(Reasoning): 쿼리를 받으면 에이전트는 직접 답할 충분한 정보가 있는지, 아니면 도구를 써야 하는지 평가합니다.
  • 행동(Acting): 도구를 쓰기로 결정하면 도구를 실행하고, 다시 추론 단계로 돌아가 이제 쿼리에 답할 수 있는지 또는 추가 도구 사용이 필요한지 판단합니다.

에이전트와 대화할 프롬프트 템플릿을 만듭니다.

# create a prompt template to chat with an agent
q_template = (
    "I am interested in {topic}. \n"
    "Find papers in your knowledge database related to this topic; use the following template to query research_paper_query_engine_tool tool: 'Provide title, summary, authors and link to download for papers related to {topic}'. If there are not, could you fetch the recent one from arXiv? \n"
)

answer = agent.chat(q_template.format(topic="Audio-Language Models"))

Markdown(answer.response)

에이전트는 RAG 도구를 선택해 관련 논문을 찾고 요약해 주었습니다. 에이전트가 채팅 기록을 유지하므로, 논문을 명시적으로 언급하지 않아도 다운로드를 요청할 수 있어요.

answer = agent.chat("Download the papers, which you mentioned above")

RAG에 없는 주제로 물어보면 어떻게 되는지 살펴봅니다.

answer = agent.chat(q_template.format(topic="Gaussian process"))

보시다시피 에이전트는 저장소에서 논문을 찾지 못해 Arxiv에서 가져왔습니다.

더 알아보기 (Learn more)