Langchain - 2차 과제

from langchain_community.document_loaders import TextLoader
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import CharacterTextSplitter
from langchain_chroma import Chroma

# 텍스트 파일을 인코딩(UTF-8)하여 로드
raw_documents = TextLoader('./state_of_the_union.txt', encoding='utf-8').load()

# 텍스트를 일정한 크기(1000자)로 나누고, 중복되는 부분이 없도록 설정
# 텍스트를 분할하고 저장
# 분할된 문서들을 임베딩하여 DB 구축
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
documents = text_splitter.split_documents(raw_documents)
db = Chroma.from_documents(documents, OpenAIEmbeddings(model="text-embedding-3-small"))

# 검색할 쿼리 설정
query = input("쿼리: ")

# 유사도 검색(similarity_search)
print("\n======= Similarity Search ======")

# DB에 저장된 문서 중 쿼리와 가장 유사한 문서 검색
# 검색된 문서들 중 첫번째 문서 내용 출력
docs = db.similarity_search(query)
print(docs[0].page_content)

# 최대 한계 관련성(Maximum Marginal Relevance, MMR) 검색
# 쿼리와 관련성이 높으면서도 서로 다양한 내용을 포함하는 문서들을 선택하는 검색 기법
# 1) 먼저 쿼리와 유사도가 높은 fetch_k개의 문서를 검색하고
# 2) 쿼리와 가장 유사한 문서를 하나 선택
# 3) 나머지 문서들 중 이미 선택된 문서와는 다른 내용을 가진 문서를 선택
# 4) 이 과정을 k개의 문서가 선택될 때까지 반복
print("\n======= MMR ======")

# 10개의 문서 검색 후 2개의 문서를 선택
mmr_docs = db.max_marginal_relevance_search(query, k=2, fetch_k=10)

# 검색 결과로 선택된 문서들을 순서대로 출력
for i in range(len(mmr_docs)):
    print(f"\n<{i+1}> {mmr_docs[i].page_content}")

from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import FAISS
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

# FAISS(Facebook AI Similarity Search)는 Facebook AI Research에 의해 개발된 라이브러리
# 대규모 벡터 데이터셋에서 유사도 검색을 빠르고 효율적으로 수행
# 특히 벡터의 압축된 표현을 사용하여 메모리 사용량을 최소화하면서도 검색 속도를 극대화

# FAISS 라이브러리로 벡터 스토어 생성
# FAISS는 주어진 텍스트 데이터 "harrison worked at kensho"에 대해 임베딩을 생성하여 저장
vectorstore = FAISS.from_texts(
    ["harrison worked at kensho"], embedding=OpenAIEmbeddings(model="text-embedding-3-small") 
)

# 벡터 스토어를 검색기로 변환, 검색기로 쿼리 관련 문서 검색 가능
retriever = vectorstore.as_retriever()

# 대화 모델 생성
model = ChatOpenAI(model="gpt-3.5-turbo")

# 주어진 문맥(context)을 기반으로 질문(question)에 답변하는 형식의 프롬프트 템플릿 정의
# 템플릿에서 프롬프트 객체 생성
template1 = """Answer the question based only on the following context:
{context}

Question: {question}
"""
prompt1 = ChatPromptTemplate.from_template(template1)

# 체인 설정
# 질문을 받아 관련 문서를 검색하고 결과를 context로 전달
retrieval_chain = (
    {"context": retriever, "question": RunnablePassthrough()} 
    | prompt1
    | model
    | StrOutputParser()
)

# 체인을 실행하여 결과 출력
answer1 = retrieval_chain.invoke("where did harrison work?")
print(answer1)

# 첫번째 체인의 응답을 이탈리아로 번역하는 프롬프트 템플릿, 체인 설정 후 실행
prompt2 = ChatPromptTemplate.from_template("Translate {answer} into Italian.")
translation_chain = prompt2 | model | StrOutputParser()
answer2 = translation_chain.invoke({"answer": answer1})
print(answer2)

# OpenAI와 SerpAI의 API 키를 환경변수로 등록함

from langchain import hub
from langchain.agents import AgentExecutor, create_react_agent
from langchain_openai import ChatOpenAI
from langchain_community.agent_toolkits.load_tools import load_tools 

# 검색 작업을 캡슐화하는 Runnable 클래스 정의
# 실행 가능한 작업 단위인 Runnable에 해당함
class SearchRunnable:   
   # 초기화 메소드, 검색 쿼리를 받아 저장
   def __init__(self, query): 
       """
       Initializes a new instance of the SearchRunnable class.
       Args: 
           query (str): The search query.
       """   
       self.query = query

   # 에이전트를 사용하여 검색을 수행하고 결과를 반환
   def run(self, agent):
       """
       Performs the search using the provided agent and returns the search result.
       Args:
           agent: The search agent to use for the search.
       Returns:
           str: The search result.
       """  
       result = agent.invoke({"input": self.query})
       return result["output"]

# 모델 생성
llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)

# SerpAPI 도구를 로드하여 검색 기능 제공
tools = load_tools(["serpapi"], llm=llm)

# "hwchase17/react" 프롬프트 템플릿을 가져오기
prompt = hub.pull("hwchase17/react")

# 모델, 도구, 프롬프트를 사용하여 React 에이전트 생성
# 생성된 에이전트와 도구를 사용하여 AgentExecutor로 검색 에이전트 생성
agent = create_react_agent(llm, tools, prompt)
search_agent = AgentExecutor(agent=agent, tools=tools, verbose=True)

# 검색 수행 함수
def search(query: str) -> str:
   """
   Performs a search using the given query and returns the search result.
   Args:
       query (str): The query string to search for.
   Returns:
       str: The search result found by the search agent.
   """
   runnable = SearchRunnable(query) # 주어진 쿼리로 SearchRunnable 객체 생성
   return runnable.run(search_agent) # SearchRunnable의 run 메소드 호출, 검색 에이전트 검색 결과 반환

# 사용자로부터 검색 쿼리를 입력받기
user_query = input("질문 입력: ")

# search 함수로 사용자 쿼리에 대한 검색 결과 얻기, 출력
result = search(user_query)
print("검색 결과:", result)