Um exemplo prático de implementação de um Retrieval-Augmented Generation (RAG)

Introdução

No mundo dinâmico da inteligência artificial, os chatbots têm desempenhado um papel fundamental na melhoria da interação entre empresas e clientes. Hoje, exploraremos uma abordagem inovadora para a criação de chatbots, combinando a potência do ChatGPT com a tecnologia LangChain.

Imagine ter um chatbot que não apenas responda às perguntas dos usuários, mas que também seja capaz de buscar informações específicas em seus próprios documentos, proporcionando respostas mais contextuais e precisas. É exatamente isso que a abordagem Retrieval-Augmented Generation (RAG) oferece. Neste artigo, vamos guiar você pelo processo de construção de um chatbot personalizado que utiliza seus próprios documentos como fonte de conhecimento.

Sobre o RAG

O Retrieval-Augmented Generation (RAG) é uma abordagem avançada na área de processamento de linguagem natural (NLP) que combina mecanismos de recuperação de informações com modelos de geração de texto. Em vez de depender exclusivamente de modelos generativos para criar respostas, o RAG integra um componente de recuperação de informações. Esse componente ajuda a extrair informações contextuais relevantes de um banco de dados e fornece esses dados ao modelo de geração de texto, possibilitando a produção de respostas mais precisas e contextualizadas.

Veja como funciona a abordagem:

1. O usuário faz uma pergunta para o modelo;
2. A pergunta é passada para o Retrieval Model;
3. O Retrieval Model recupera da base de dados os documentos relevantes para responder a pergunta do usuário;
4. O Retrieval Model envia ao LLM um prompt contendo a pergunta do usuário e as informações relevantes presentes nos documentos recuperados;
5. A LLM pré-treinada gera uma resposta baseada nas informações fornecidas e a retorna ao usuário.

Implementação

Sem mais delongas, vamos iniciar preparando nosso ambiente virtual com os pré-requisitos necessários para o funcionamento da nossa aplicação.

Observação: Neste guia vamos assumir que o leitor já possui conhecimentos prévios de programação em Python.

Requisitos

Vamos criar uma pasta chamada chatbot e iniciar um ambiente virtual:

‍
python3 -m venv venv

‍

Ative o ambiente virtual e instale as dependências da aplicação:

‍

source venv/bin/activate

pip install --upgrade pip

pip install python-dotenv langchain langchain-openai openai milvus pymilvus unstructured tiktoken lark

‍

Na pasta raiz do projeto, crie um arquivo .env e adicione as seguintes variáveis:

‍

OPENAI_API_KEY = ""
MILVUS_HOST = "localhost"
MILVUS_PORT = "19530"

‍

Substitua o valor de OPENAI_API_KEY pela sua chave de API da OpenAI.

Caso queira usar um servidor Milvus em outra localidade, sinta-se à vontade para substituir as variáveis MILVUS_HOST e MILVUS_PORT.

Modelo

Agora vamos implementar nosso RAG. Comece criando um arquivo chamado model.py na pasta raiz.

Abra o arquivo usando seu editor favorito e importe as dependências necessárias:

‍

import os
from dotenv import load_dotenv
load_dotenv()

from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain.chains.query_constructor.base import AttributeInfo
from langchain.prompts import ChatPromptTemplate
from langchain.schema.output_parser import StrOutputParser
from langchain.memory import ConversationTokenBufferMemory
from langchain_core.prompts import MessagesPlaceholder
from langchain_openai.embeddings import OpenAIEmbeddings
from langchain_openai.chat_models import ChatOpenAI
from langchain_community.document_loaders import DirectoryLoader
from langchain_community.vectorstores import Milvus
from milvus import default_server as milvus_server

‍

Agora vamos criar nossa classe RAG, onde vamos implementar toda a lógica do nosso modelo de recuperação de dados:

‍

class RAG():
   def __init__(self,
                docs_dir: str,
                n_retrievals: int = 4,
                chat_max_tokens: int = 3097,
                model_name = "gpt-3.5-turbo",
                creativeness: float = 0.7):
       self.__model = self.__set_llm_model(model_name, creativeness)
       self.__docs_list = self.__get_docs_list(docs_dir)
       self.__retriever = self.__set_retriever(k=n_retrievals)
       self.__chat_history = self.__set_chat_history(max_token_limit=chat_max_tokens)

‍

Como podem ver, nossa classe recebe 5 parâmetros, sendo apenas 1 obrigatório (uma string contendo o nome da pasta de documentos). Além disso, ao ser instanciado, nosso objeto RAG atribui valores a 4 variáveis privadas:

1. __model: o objeto do nosso modelo GPT da OpenAI;
2. __docs_list: a lista de documentos carregados;
3. __retriever: o retriever que usamos para recuperar os dados;
4. __chat_history: o buffer que usamos para armazenar em memória o histórico de conversa do nosso chat.

Métodos privados

Agora vamos implementar os nossos 4 métodos privados que são usados justamente para atribuir valores às nossas variáveis privadas mencionadas acima. Vamos lá:

1. LLM Model:

Vamos instanciar o modelo gpt-3.5-turbo usando a classe ChatOpenAI fornecida pelo LangChain. Aqui é importante salientar que, apesar de estarmos usando o modelo da OpenAI, o LangChain também tem integração com diversos outros modelos de GenAI. Para mais informações sobre outros modelos acesse a documentação oficial.

def __set_llm_model(self, model_name = "gpt-3.5-turbo", temperature: float = 0.7):
       return ChatOpenAI(model_name=model_name, temperature=temperature)

‍

2. Docs List:

Para ler nossos documentos, vamos usar o DirectoryLoader do LangChain. Note que habilitamos o modo recursivo (para ler subpastas) e o modo multithreading (para executar paralelamente em mais de um núcleo do processador). No nosso exemplo estamos usando um máximo de 4 núcleos, mas é possível mudar isso alterando a variável max_concurrency.

Usamos a função load_and_split() para carregar nossos arquivos. Essa função automaticamente divide grandes arquivos em outros arquivos menores, garantindo a leitura e armazenagem correta para o nosso modelo.

‍

def __get_docs_list(self, docs_dir: str) -> list:
       print("Carregando documentos...")
       loader = DirectoryLoader(docs_dir,
                                recursive=True,
                                show_progress=True,
                                use_multithreading=True,
                                max_concurrency=4)
       docs_list = loader.load_and_split()
     
       return docs_list

‍

3. Retriever:

Antes de partir para o retriever, precisamos criar nosso vector store. Para isso estamos usando o banco de dados vetorial Milvus. Note que em collection_name damos um nome descritivo para a nossa coleção de dados. Caso queira, fique à vontade para modificá-lo.

Agora que temos o vector store, podemos iniciar a criação do nosso retriever. Primeiro precisamos criar a variável metadata_field_info, que armazena informações sobre os metadados dos nossos documentos. Por padrão, ao ler os documentos, o loader adiciona o metadado “source”, onde armazena o caminho e nome do arquivo ao qual se refere. Aqui então, apenas damos mais detalhes sobre esse metadado para que o modelo consiga interpretar do que se trata. Caso você adicione mais metadados aos seus documentos, não esqueça de descrevê-los nesta lista. Além disso, também criamos a variável document_content_description para dizer ao modelo sobre o que são nossos documentos.

Agora que temos nosso vector store e nossas variáveis de metadados, podemos criar nosso retriever propriamente dito. Para esse exemplo, vamos usar o Self-Querying Retriever, passando como parâmetros o nosso modelo de IA, o vector store, nossas informações de metadados e um último parâmetro “k”. O parâmetro “k” fornece o valor máximo de documentos que devem ser recuperados pelo retriever a cada execução.

def __set_retriever(self, k: int = 4):
       # Milvus Vector Store
       embeddings = OpenAIEmbeddings()
       milvus_server.start()
       vector_store = Milvus.from_documents(
           self.__docs_list,
           embedding=embeddings,
           connection_args={"host": os.getenv("MILVUS_HOST"), "port": os.getenv("MILVUS_PORT")},
           collection_name="personal_documents",
       )

       # Self-Querying Retriever
       metadata_field_info = [
           AttributeInfo(
               name="source",
               description="O caminho de diretórios onde se encontra o documento",
               type="string",
           ),
       ]

       document_content_description = "Documentos pessoais"

       _retriever = SelfQueryRetriever.from_llm(
           self.__model,
           vector_store,
           document_content_description,
           metadata_field_info,
           search_kwargs={"k": k}
       )

       return _retriever

‍

4. Chat History:

Por último, usamos um Conversation Token Buffer para armazenar em memória o histórico da conversa. Como estamos usando o modelo gpt-3.5-turbo da OpenAI, que possui um limite de apenas 4097 tokens por requisição, vamos usar este buffer por ser capaz de armazenar os dados até um limite pré-definido de tokens (no nosso caso, 3097). Quando o buffer chega em seu limite de tokens, automaticamente ele começa a descartar mensagens mais antigas, dando prioridade à memória para mensagens mais recentes.

def __set_chat_history(self, max_token_limit: int = 3097):
       return ConversationTokenBufferMemory(llm=self.__model, max_token_limit=max_token_limit, return_messages=True)

Métodos públicos

Agora que já temos nossos métodos privados e toda a lógica de inicialização de nossas variáveis privadas no nosso construtor, vamos criar nosso método público ask(). Este será o método responsável por receber uma pergunta do usuário e retornar uma resposta apropriada. Vamos lá:

def ask(self, question: str) -> str:
       prompt = ChatPromptTemplate.from_messages([
           ("system", "Você é um assistente responsável por responder perguntas sobre documentos. Responda a pergunta do usuário com um nível de detalhes razoável e baseando-se no(s) seguinte(s) documento(s) de contexto:\n\n{context}"),
           MessagesPlaceholder(variable_name="chat_history"),
           ("user", "{input}"),
       ])
     
       output_parser = StrOutputParser()
       chain = prompt | self.__model | output_parser
       answer = chain.invoke({
           "input": question,
           "chat_history": self.__chat_history.load_memory_variables({})['history'],
           "context": self.__retriever.get_relevant_documents(question)
       })

       # Atualização do histórico de conversa
       self.__chat_history.save_context({"input": question}, {"output": answer})
     
       return answer

‍

Primeiramente, criamos um ChatPromptTemplate a partir de uma lista de mensagens, que traz o histórico de mensagens salvo no buffer através do objeto MessagesPlaceholder. Tendo nosso template, criamos nossa chain e a invocamos passando a variável com a pergunta do usuário, as mensagens contidas no histórico e os documentos relevantes para a resposta da pergunta (obtidos através da função get_relevant_documents). Por fim, atualizamos o histórico de perguntas e retornamos a resposta fornecida pelo modelo de GenAI.

Execução

Pronto! Já temos nosso RAG devidamente implementado! Agora precisamos criar o arquivo que será responsável pela execução da nossa aplicação. Na pasta raiz do projeto, crie um arquivo chamado main.py e adicione o seguinte:

‍

from model import RAG

rag = RAG(
   docs_dir='docs', # Nome do diretório onde estão os documentos
   n_retrievals=1, # Número de documentos retornados pela busca (int)  :   default=4
   chat_max_tokens=3097, # Número máximo de tokens que podem ser usados na memória do chat (int)  :   default=3097
   creativeness=1.2, # Quão criativa será a resposta (float 0-2)  :   default=0.7
)

print("\nDigite 'sair' para sair do programa.")
while True:
   question = str(input("Pergunta: "))
   if question == "sair":
       break
   answer = rag.ask(question)
 print('Resposta:', answer)

‍

Aqui nós importamos nosso objeto RAG criado no arquivo model.py e instanciamos com os parâmetros desejados (fique à vontade para alterá-los como queira). Depois disso, usamos um laço de repetição while para manter nosso chat em execução, esperando uma pergunta do usuário.

Como podem notar, nós passamos “docs” como valor para o parâmetro docs_dir, ou seja, devemos criar na raiz do projeto uma pasta chamada docs e adicionar todos nossos documentos dentro dela. Nosso RAG então irá ler os documentos contidos nesta pasta, encontrar quais são relevantes para responder a pergunta do usuário e usar o GPT-3.5 para criar uma resposta apropriada com os dados fornecidos.

E voilà! Agora que temos nosso chatbot, apenas execute o seguinte comando no terminal:

‍

python main.py

‍

e divirta-se conversando com ele sobre seus documentos!

Próximos passos

Para aprimorar ainda mais a aplicação, sugerimos construir uma interface de usuário para tornar o chat mais intuitivo e dinâmico. Com o objetivo de facilitar o desenvolvimento, recomendamos as seguintes bibliotecas:

• LangServe: o próprio LangChain fornece esta biblioteca. Ela tem um foco maior na criação de REST APIs, mas também acompanha uma interface de usuário integrada simples para configurar e rodar o aplicativo com saída de streaming e visibilidade nas etapas intermediárias. Uma opção mais simples e rápida para quem não tem muito interesse em personalizar a interface.
• Streamlit: Uma biblioteca Python eficiente para criar interfaces de usuário interativas e atrativas, simplificando o desenvolvimento de aplicativos web sem a necessidade de habilidades avançadas em design ou programação front-end.
• Dash: Uma biblioteca Python para construção de aplicativos analíticos interativos com o uso de componentes web reutilizáveis.

Repositório

Todo o código desenvolvido para a produção deste artigo pode ser encontrado no seguinte repositório do GitHub: DP6/rag-chatbot