Como usar o AutoGen para construir um agente de IA

Neste tutorial, você vai aprender a usar o framework AutoGen para construir um agente capaz de responder perguntas com base em documentos armazenados em um banco de dados vetorial.

O que é o AutoGen?

O AutoGen é o resultado de uma pesquisa colaborativa entre a Microsoft, a Universidade Estadual da Pensilvânia (Penn State University) e a Universidade de Washington. É um framework projetado para simplificar o gerenciamento, a melhoria e a automação dos fluxos de trabalho de Large Language Models (LLM). Ele permite criar agentes personalizáveis e conversacionais que otimizam o uso de LLMs avançados como o GPT-4, complementando suas limitações com a integração de humanos e outras ferramentas. Além disso, facilita a interação e a comunicação automática entre vários agentes. Mais informações: AutoGen | AutoGen (microsoft.github.io)

Configurar credenciais e ambiente

Primeiro, é preciso configurar as credenciais e as variáveis de ambiente necessárias para acessar os serviços do Azure.

Para isso, crie um arquivo .env no mesmo diretório do seu script com as seguintes informações:

AZURE_TENANT_ID = "00000000-0000-0000-0000-000000000000"
AZURE_CLIENT_ID = "00000000-0000-0000-0000-000000000000"
AZURE_CLIENT_SECRET = "xxxxx"

Inicialização e configuração

O primeiro passo é inicializar as credenciais usando o ChainedTokenCredential e o EnvironmentCredential do Azure. Depois, obtemos um token de acesso para os serviços cognitivos do Azure.

import os
from azure.identity import ChainedTokenCredential, EnvironmentCredential
from dotenv import load_dotenv

load_dotenv()
credential = ChainedTokenCredential(EnvironmentCredential())
access_token = credential.get_token("https://cognitiveservices.azure.com/.default")

Configuração do modelo e das variáveis de ambiente

Configuramos o modelo e as variáveis de ambiente necessárias para usar a OpenAI no Azure.

deployment = "zero"
embedding_deployment = "embedding-ada-zero"

os.environ["AZURE_OPENAI_ENDPOINT"] = "https://zerogap.openai.azure.com/"
os.environ["AZURE_OPENAI_API_KEY"] = access_token.token
os.environ["OPENAI_API_TYPE"] = "azure_ad"
os.environ["OPENAI_DEPLOYMENT"] = deployment

Inicialização do modelo e dos embeddings

Inicializamos o modelo AzureChatOpenAI e a classe AzureOpenAIEmbeddings.

from langchain_openai import AzureOpenAIEmbeddings, AzureChatOpenAI

llm = AzureChatOpenAI(openai_api_version="2023-07-01-preview", azure_deployment=deployment, temperature=0.5)

embeddings = AzureOpenAIEmbeddings(
    azure_deployment=embedding_deployment,
    openai_api_version="2023-07-01-preview",
    chunk_size=1
)

Configuração do agente AutoGen

Configuramos o agente AutoGen com as informações necessárias para se conectar aos serviços do Azure.

config_list = [
  {
    "model": deployment,
    "api_type": "azure",
    "api_key": os.environ['AZURE_OPENAI_API_KEY'],
    "base_url": os.environ["AZURE_OPENAI_ENDPOINT"],
    "api_version": "2024-02-01"
  }
]

Carregar documentos e criar o retriever

Carregamos os documentos de um diretório persistente e configuramos o retriever para buscar no banco de dados vetorial.

from langchain_community.vectorstores import Chroma

persist_directory = "chroma_db_generic"
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
vectordb.get()

retriever = vectordb.as_retriever(search_type="similarity", search_kwargs={"k": 5})

Formatar documentos

Criamos uma função para formatar os documentos recuperados.

def format_docs(docs):
    if docs:
        return "\n\n".join(doc.page_content for doc in docs)

Ferramenta: obter documentos

Definimos uma função que o agente vai usar para responder perguntas com base nos documentos recuperados.

def get_documents(question: str) -> str:
    context = retriever.invoke(question)
    formated_context = format_docs(context)
    return formated_context

Configuração dos agentes de conversa

Configuramos os agentes de conversa: user_proxy e assistant.

from autogen import register_function, ConversableAgent

user_proxy = ConversableAgent(
    name="User",
    is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
    human_input_mode="NEVER",
    max_consecutive_auto_reply=10
)

assistant = ConversableAgent(
    name="Assistant",
    system_message="""You are the ZeroGap AI Bot. Your role is to provide answers strictly based on the Zerogap Documents. Keep answers as short as possible. Follow these instructions step by step:
    1. Answer the user's question using any tool available.
    2. If the context is not clear. Give the user a message to clarify the context. If possible, offer some options from the retrieved context. Do NOT answer the question.
    3. Return 'TERMINATE' when the task is done.""",
    llm_config={"config_list": config_list, "cache_seed": None},
)

Registrar a ferramenta e executar o agente

Registramos a ferramenta get_documents e definimos a função para executar o agente com uma consulta.

register_function(
    get_documents,
    caller=assistant,
    executor=user_proxy,
    name="retrieve_tool",
    description="Gets information from Zerogap Documents",
)

def run_agent(query):
    chat_result = user_proxy.initiate_chat(
        assistant,
        message=query,
        summary_method="reflection_with_llm",
        summary_args={"summary_prompt" : "Return the final response to the user's query. Do not include the user's query in the response. Do not include the system message in the response."},
        max_turns=2
    )
    return chat_result.summary

Execução do chatbot

Por fim, executamos o loop principal do chatbot.

print("\033[92m" + "Welcome to the ZEROGAP AI QA Bot!" + "\033[0m")
print("\033[92m" + "###############################" + "\033[0m")
print("\033[92m" + "###############################" + "\033[0m")
print("\n")
while True:
    question = input("\033[93m" + "You: " + "\033[0m")
    print("\n")
    if question == "quit":
        break
    result = run_agent(question)
    print("\033[92m" + result + "\033[0m")
    print("\n")

Código completo

import os
from azure.identity import ChainedTokenCredential, EnvironmentCredential
from langchain_community.vectorstores import Chroma
from langchain_openai import AzureOpenAIEmbeddings
from langchain_openai import AzureChatOpenAI
from autogen import register_function, ConversableAgent
from dotenv import load_dotenv
load_dotenv()

# Place a .env file within the same folder with the following information:
# AZURE_TENANT_ID = "00000000-0000-0000-0000-000000000000"
# AZURE_CLIENT_ID = "00000000-0000-0000-0000-000000000000"
# AZURE_CLIENT_SECRET = "xxxxx"

credential = ChainedTokenCredential(EnvironmentCredential())
access_token = credential.get_token("https://cognitiveservices.azure.com/.default")

# Model
deployment = "zerogap"
# Model text-embedding-ada-002
embedding_deployment = "embedding-ada-zero"

# Set OS environment variables
os.environ["AZURE_OPENAI_ENDPOINT"] = "https://zerogap.openai.azure.com/"
os.environ["AZURE_OPENAI_API_KEY"] = access_token.token
os.environ["OPENAI_API_TYPE"] = "azure_ad"
os.environ["OPENAI_DEPLOYMENT"] = deployment

llm = AzureChatOpenAI(openai_api_version="2023-07-01-preview", azure_deployment=deployment, temperature=0.5)

# Initialize the OpenAIEmbeddings class
embeddings = AzureOpenAIEmbeddings(
    azure_deployment=embedding_deployment,
    openai_api_version="2023-07-01-preview",
    chunk_size=1
)

# Set the configuration for AutoGen Agent
config_list = [
  {
    "model": deployment,
    "api_type": "azure",
    "api_key": os.environ['AZURE_OPENAI_API_KEY'],
    "base_url": os.environ["AZURE_OPENAI_ENDPOINT"],
    "api_version": "2024-02-01"
  }
]

# Load documents from the persisted directory
# Assuming you already have the Document loaded to a vector database
persist_directory = "chroma_db_generic"
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
vectordb.get()

retriever = vectordb.as_retriever(search_type="similarity", search_kwargs={"k": 5})

# Join all the documents from the retriever together with newlines
def format_docs(docs):
    if docs:
        return "\n\n".join(doc.page_content for doc in docs)

# Tool - the function that the Agent will use to answer the questions
def get_documents(question: str) -> str:
    # Queries the Vector DB using user's question
    context = retriever.invoke(question)
    formated_context = format_docs(context)
    return formated_context

user_proxy = ConversableAgent(
    name="User",
    is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
    human_input_mode="NEVER",
    max_consecutive_auto_reply=10
)

assistant = ConversableAgent(
    name="Assistant",
    system_message="""You are the ZeroGap AI Bot. Your role is to provide answers strictly based on the Zerogap Documents. Keep answers as short as possible. Follow these instructions step by step:
    1. Answer the user's question using any tool available.
    2. If the context is not clear. Give the user a message to clarify the context. If possible, offer some options from the retrieved context. Do NOT answer the question.
    3. Return 'TERMINATE' when the task is done.""",
    llm_config={"config_list": config_list, "cache_seed": None},
)

# Register the Document tool.
register_function(
    get_documents,
    caller=assistant,
    executor=user_proxy,
    name="retrieve_tool",
    description="Gets information from Zerogap Documents",
)

# Run the agent
def run_agent(query):
    chat_result = user_proxy.initiate_chat(
        assistant,
        message=query,
        summary_method="reflection_with_llm",
        summary_args={"summary_prompt" : "Return the final response to the user's query. Do not include the user's query in the response. Do not include the system message in the response."},
        max_turns=2
    )
    return chat_result.summary

print("\033[92m" + "Welcome to the ZEROGAP AI QA Bot!" + "\033[0m")
print("\033[92m" + "###############################" + "\033[0m")
print("\033[92m" + "###############################" + "\033[0m")
print("\n")
while True:
    question = input("\033[93m" + "You: " + "\033[0m")
    print("\n")
    if question == "quit":
        break
    result = run_agent(question)
    print("\033[92m" + result + "\033[0m")
    print("\n")