Cómo usar AutoGen para construir un Agente de IA

En este tutorial, aprenderás a usar el Framework AutoGen para construir un agente que pueda responder preguntas basadas en documentos almacenados en una base de datos vectorial.

¿Qué es AutoGen?

AutoGen es el resultado de una investigación colaborativa entre Microsoft, la Universidad Estatal de Pensilvania (Penn State University) y la Universidad de Washington. Es un Framework diseñado para simplificar la gestión, mejora y automatización de los flujos de trabajo de Modelos de Lenguaje a Gran Escala (LLM). Permite la creación de agentes personalizables y conversacionales que optimizan el uso de LLM avanzados como GPT-4, complementando sus limitaciones con la integración de humanos y otras herramientas. Facilita además la interacción y comunicación automática entre múltiples agentes. Más información: AutoGen | AutoGen (microsoft.github.io)

Configurar Credenciales y Entorno

Primero, es necesario configurar las credenciales y variables de entorno necesarias para acceder a los servicios de Azure.

Para ello, crea un archivo .env en el mismo directorio que tu script con la siguiente información:

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

Inicialización y Configuración

El primer paso es inicializar las credenciales utilizando el ChainedTokenCredential y EnvironmentCredential de Azure. Luego, obtenemos un token de acceso para los servicios cognitivos de 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")

Configuración del Modelo y Variables de Entorno

Configuramos el modelo y las variables de entorno necesarias para el uso de OpenAI en 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

Inicialización del Modelo y Embeddings

Inicializamos el modelo AzureChatOpenAI y la clase 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
)

Configuración del Agente AutoGen

Configuramos el agente AutoGen con la información necesaria para conectar con los servicios de 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"
  }
]

Cargar Documentos y Crear el Recuperador

Cargamos los documentos desde un directorio persistente y configuramos el recuperador para buscar en la base de datos vectorial.

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})

Formatear Documentos

Creamos una función para formatear los documentos recuperados.

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

Herramienta: Obtener Documentos

Definimos una función que el agente usará para responder preguntas basándose en los documentos recuperados.

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

Configuración de los Agentes de Conversación

Configuramos los agentes de conversación: user_proxy y 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 la Herramienta y Ejecutar el Agente

Registramos la herramienta get_documents y definimos la función para ejecutar el agente con una 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

Ejecución del Chat Bot

Finalmente, ejecutamos el bucle principal del 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")