How to use AutoGen to build an AI Agent

Translated from the Spanish original. Read in Spanish

In this tutorial, you’ll learn how to use the AutoGen Framework to build an agent that can answer questions based on documents stored in a vector database.

iStock AI Generator

What is AutoGen?

AutoGen is the result of joint research by Microsoft, Pennsylvania State University (Penn State) and the University of Washington. It’s a framework designed to simplify the management, improvement and automation of Large Language Model (LLM) workflows. It enables customisable, conversational agents that make the most of advanced LLMs such as GPT-4, making up for their limitations by bringing in humans and other tools. It also makes automatic interaction and communication between multiple agents easier. More information: AutoGen | AutoGen (microsoft.github.io)

Set Up Credentials and Environment

First, you need to set up the credentials and environment variables required to access the Azure services.

To do so, create a .env file in the same directory as your script with the following information:

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

Initialisation and Configuration

The first step is to initialise the credentials using Azure’s ChainedTokenCredential and EnvironmentCredential. Then we get an access token for Azure Cognitive Services.

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

Configuring the Model and Environment Variables

We configure the model and the environment variables needed to use OpenAI on 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

Initialising the Model and Embeddings

We initialise the AzureChatOpenAI model and the AzureOpenAIEmbeddings class.

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
)

Configuring the AutoGen Agent

We configure the AutoGen agent with the information needed to connect to the Azure services.

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 and Create the Retriever

We load the documents from a persistent directory and set up the retriever to search the vector database.

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

Format Documents

We create a function to format the retrieved documents.

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

Tool: Get Documents

We define a function the agent will use to answer questions based on the retrieved documents.

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

Configuring the Conversational Agents

We configure the conversational agents: user_proxy and 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},
)

Register the Tool and Run the Agent

We register the get_documents tool and define the function that runs the agent with a query.

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

Running the Chat Bot

Finally, we run the chatbot’s main loop.

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

Complete Code

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

Maximiliano Díaz Doglia

AI Platform Engineer & Full-Stack Developer
Building Enterprise Integrations & Automations

Published in: AI