TavilySearchAPI를 이용한 Retriever(v.1)

langchain에서 사용되는 다양한 Tool 중에 LLM의 AI agent를 위해 특별히 제작된 검색 엔진으로, 정확하고 사실에 입각한 실시간 검색 결과를 빠르게 제공한다고 알려져 있는 “Tavily’s Search API”를 retriever로 활용하는 예제를 연습해 봤습니다.

Tavily’s Search API를 사용하기 전에 2가지 준비 과정이 필요합니다.

Tavily API Key 발급
tavily-python package 설치

1. Tavily API Key

Tavily API Key는 Tavily 홈페이지에 가입해서 발급 받으면 됩니다.

가입하면 “Researcher” plan이 적용되는데, 월 1,000건의 API calls이 제공된다고 하네요.

전체 Taviliy의 요금제는 다음과 같습니다. (https://tavily.com/#pricing)

이렇게 발급한 API Key를 시스템 환경 변수에 등록하여 사용하였습니다.

2. tavily-python package 설치

pip install -U langchain-community tavily-python

활용 예제

“Explain about Euro 2024”이라는 질문에 대해 Tavily Search API로 관련 정보를 검색(retrieve)하고, 그 정보를 활용하여 LLM이 답변을 생성한다.

작성 코드

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_community.retrievers import TavilySearchAPIRetriever
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
import os

llm = ChatOpenAI()

# 검색한 문서 결과를 하나의 문단으로 합친다.
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

# 시스템 환경 변수 TAVILY API KEY 가져오기
TAVILY_API_KEY = os.environ.get('TAVILY_API_KEY')

retriever = TavilySearchAPIRetriever(k = 3, search_depth="advanced")
# default of k : 10, search_depth : basic

# Question
input_query = "Explain about Euro 2024."
input_query1 = "What is the timetable for Euro 2024?"

# retriever 결과에 대한 출력
documents = retriever.invoke(input_query)

for idx, doc in enumerate(documents):
    print(f"[{idx+1}]")
    page_content = doc.page_content
    metadata = doc.metadata
    
    title = metadata['title']
    source = metadata['source']
    score = metadata['score']
    images = metadata['images']
    
    print(f"* page_content: {page_content}")
    print(f"* title: {title}")
    print(f"* source: {source}")
    print(f"* score: {score}")
    print(f"* images: {images}")
    print('-'*20)

prompt = ChatPromptTemplate.from_template("""Answer the following question based only on the provided context:

<context>
{context}
</context>

Question: {input}
""")

runnable = RunnablePassthrough.assign(
    context=(lambda x: x['input']) 
    | retriever 
    | RunnableLambda(format_docs))

# chaining

chain = (
    runnable
    | prompt 
    | llm 
    | StrOutputParser()
)

response = chain.invoke({"input": input_query})

print('*'*10, "답변 출력", '*'*10)
print(response)

참고 사항

TavilySearchAPIRetriever()
- k : 검색되는 결과의 수(default 10)
- search_depth : basic(quick response) or advanced(in-depth, quality results)

Retriever 결과

[1]
* page_content: Article top media content
Article body
Where will EURO 2024 be held?
Germany will host EURO 2024, having been chosen to stage the 17th edition of the UEFA European Championship at a UEFA Executive Committee meeting in Nyon on 27 September 2018. Host cities
EURO 2024 fixtures by venue
EURO 2024 fixtures by team
Also visit
Change language
Services links and disclaimer
© 1998-2024 UEFA. Where and when will the final of UEFA EURO 2024 be played?
Berlin's Olympiastadion will stage the final on Sunday 14 July 2024.
 The ten venues chosen to host games at the tournament include nine of the stadiums used at the 2006 World Cup plus the Düsseldorf Arena.
 All you need to know
Thursday, January 11, 2024
Article summary
Three-time winners Germany will stage the UEFA European Championship in 2024.

* title: EURO 2024: All you need to know | UEFA EURO 2024
* source: https://www.uefa.com/euro2024/news/0257-0e13b161b2e8-4a3fd5615e0c-1000--euro-2024-all-you-need-to-know/
* score: 0.91957
* images: None
--------------------
[2]
* page_content: The EURO 2024 tournament in Germany will take center stage this summer and all eyes will be on which European giant can win it all. [ MORE: Schedule for EURO 2024] From how it works to where the games will be played and who to keep a close eye on, below is everything you need to know on the competition.
* title: UEFA EURO 2024: How does it work, qualifying, top players, favorite ...
* source: https://www.nbcsports.com/soccer/news/uefa-euro-2024-how-does-it-work-qualifying-top-players-favorite-host-city-final
* score: 0.91043
* images: None
--------------------
[3]
* page_content: 4 of 6 |. FILE - Aerial view of the Munich Euro 2024 stadium in Munich, Germany, Friday, June 7, 2024. Munich will host six matches at the European soccer Championships. The Euros kick off in Munich, Friday June 14, when host country Germany plays Scotland at Bayern Munich's Allianz Arena. The tournament begins with six groups of four teams.
* title: Euro 2024: What to know about the European Championship
* source: https://apnews.com/article/euro-2024-guide-ronaldo-4dddabf37f8073e2f9cf364a46f27e95
* score: 0.90999
* images: None
--------------------

생성된 답변

Euro 2024 is the 17th edition of the UEFA European Championship, which will be hosted by Germany. The final match of the tournament will be held at Berlin's Olympiastadion on Sunday, July 14, 2024. The tournament will feature ten venues, including nine stadiums used during the 2006 World Cup and the Düsseldorf Arena. Germany, a three-time winner of the UEFA European Championship, will be hosting the competition in 2024. The tournament will kick off in Munich on June 14, 2024, with the host country Germany playing against Scotland at Bayern Munich's Allianz Arena.

#11기 랭체인

🚀 AI스터디 16기 모집 중 (~5/8)