Author: Robert Jambrecic
Introduction
In this post, we’ll build upon the concepts introduced in our previous blog on Tools with Dependency Injection. We’ll take a deeper look at how ChatContext
can be used to manage the flow of conversations in a more structured and secure way.
By using ChatContext
, we can track and control the sequence of function calls during a conversation. This is particularly useful in situations where one task must be completed before another — for example, ensuring that a user logs in before they can check their account balance. This approach helps to prevent errors and enhances the security of the system.
Benefits of Using ChatContext
:
Flow Control: Ensures tasks are performed in the correct order, reducing the chance of mistakes.
Enhanced Security: Prevents unauthorized actions, such as accessing sensitive data before authentication.
Simplified Debugging: Logs the conversation history, making it easier to trace and resolve issues.
This blog builds on the concepts shared in the notebook.
Installation
To install AG2
, simply run the following command:
pip install ag2
Imports
import os
from typing import Annotated, Literal
from pydantic import BaseModel
from autogen.agentchat import AssistantAgent, UserProxyAgent
from autogen.tools.dependency_injection import BaseContext, ChatContext, Depends
Account and Helper Functions
The following Account
class and helper functions are adapted from the Tools with Dependency Injection post. They define the structure for securely handling account data and operations like login and balance retrieval.
Define Account Class
class Account(BaseContext, BaseModel):
username: str
password: str
currency: Literal["USD", "EUR"] = "USD"
alice_account = Account(username="alice", password="password123")
bob_account = Account(username="bob", password="password456")
account_ballace_dict = {
(alice_account.username, alice_account.password): 300,
(bob_account.username, bob_account.password): 200,
}
Helper Functions
These functions validate account credentials and retrieve account balances.
def _verify_account(account: Account):
if (account.username, account.password) not in account_ballace_dict:
raise ValueError("Invalid username or password")
def _get_balance(account: Account):
_verify_account(account)
return f"Your balance is {account_ballace_dict[(account.username, account.password)]}{account.currency}"
Agent Configuration
Configure the agents for the interaction.
config_list
defines the LLM configurations, including the model and API key.UserProxyAgent
simulates user inputs without requiring actual human interaction (set toNEVER
).AssistantAgent
represents the AI agent, configured with the LLM settings.
config_list = [{"model": "gpt-4o-mini", "api_key": os.environ["OPENAI_API_KEY"]}]
agent = AssistantAgent(
name="agent",
llm_config={"config_list": config_list},
)
user_proxy = UserProxyAgent(
name="user_proxy_1",
human_input_mode="NEVER",
llm_config=False,
)
Injecting a ChatContext Parameter
Now let’s upgrade the example from the previous post by introducing the ChatContext
parameter. This enhancement allows us to enforce proper execution order in the workflow, ensuring that users log in before accessing sensitive data like account balances.
The following functions will be registered:
login
: Verifies the user’s credentials and ensures they are logged in.get_balance
: Retrieves the account balance but only if the user has successfully logged in first.
@user_proxy.register_for_execution()
@agent.register_for_llm(description="Login")
def login(
account: Annotated[Account, Depends(bob_account)],
) -> str:
_verify_account(account)
return "You are logged in"
@user_proxy.register_for_execution()
@agent.register_for_llm(description="Get balance")
def get_balance(
account: Annotated[Account, Depends(bob_account)],
chat_context: ChatContext,
) -> str:
_verify_account(account)
# Extract the list of messages exchanged with the first agent in the conversation.
# The chat_context.chat_messages is a dictionary where keys are agents (objects)
# and values are lists of message objects. We take the first value (messages of the first agent).
messages_with_first_agent = list(chat_context.chat_messages.values())[0]
login_function_called = False
for message in messages_with_first_agent:
if "tool_calls" in message and message["tool_calls"][0]["function"]["name"] == "login":
login_function_called = True
break
if not login_function_called:
raise ValueError("Please login first")
balance = _get_balance(account)
return balance
Finally, we initiate a chat to retrieve the balance.
user_proxy.initiate_chat(agent, message="Get users balance", max_turns=4)
When user_proxy.initiate_chat(agent, message="Get users balance", max_turns=4)
is called, the following happens:
User Request: The user sends the message
"Get users balance"
, prompting the agent to suggest callingget_balance
.First Function Call: The agent attempts to execute
get_balance
, but since the user isn’t logged in, it raises an error:"Please login first"
. The agent then suggests callinglogin
to authenticate the user.Login Execution: The
login
function is executed, confirming the user’s login with"You are logged in"
.Second Function Call: With the user logged in, the agent successfully calls
get_balance
and retrieves the balance.Final Response: The agent responds with
"Your balance is 200 USD"
, confirming the correct execution of the workflow.
user_proxy_1 (to agent):
Get users balance
--------------------------------------------------------------------------------
agent (to user_proxy_1):
***** Suggested tool call (call_aeSi4R4Mo7f6JpSLTpll6lmH): get_balance *****
Arguments:
{}
****************************************************************************
--------------------------------------------------------------------------------
>>>>>>>> EXECUTING FUNCTION get_balance...
Call ID: call_aeSi4R4Mo7f6JpSLTpll6lmH
Input arguments: {}
user_proxy_1 (to agent):
***** Response from calling tool (call_aeSi4R4Mo7f6JpSLTpll6lmH) *****
Error: Please login first
**********************************************************************
--------------------------------------------------------------------------------
agent (to user_proxy_1):
***** Suggested tool call (call_7hY8yZWjrD7Irh1mKvFJL40y): login *****
Arguments:
{}
**********************************************************************
--------------------------------------------------------------------------------
>>>>>>>> EXECUTING FUNCTION login...
Call ID: call_7hY8yZWjrD7Irh1mKvFJL40y
Input arguments: {}
user_proxy_1 (to agent):
***** Response from calling tool (call_7hY8yZWjrD7Irh1mKvFJL40y) *****
You are logged in
**********************************************************************
--------------------------------------------------------------------------------
agent (to user_proxy_1):
***** Suggested tool call (call_yTSqe4kxUbCH6Kxub4wiYlIY): get_balance *****
Arguments:
{}
****************************************************************************
--------------------------------------------------------------------------------
>>>>>>>> EXECUTING FUNCTION get_balance...
Call ID: call_yTSqe4kxUbCH6Kxub4wiYlIY
Input arguments: {}
user_proxy_1 (to agent):
***** Response from calling tool (call_yTSqe4kxUbCH6Kxub4wiYlIY) *****
Your balance is 200USD
**********************************************************************
--------------------------------------------------------------------------------
agent (to user_proxy_1):
Your balance is 200 USD.
TERMINATE
--------------------------------------------------------------------------------
Conclusion
In this post, we introduced ChatContext
as an extension of dependency injection, showing how it simplifies and secures workflows in conversational AI by ensuring tasks like login happen before accessing sensitive data, improving both functionality and security.
Highlights:
Flow Control:
ChatContext
can ensure that functions are executed in the correct sequence, reducing the risk of logical errors.Enhanced Security: Sensitive actions like balance retrieval are protected by preconditions like login verification.
Simplified Debugging: With
ChatContext
, developers can trace conversations, making it easier to identify and resolve issues.
By leveraging tools like ChatContext
and employing dependency injection principles, you can build robust, secure, and user-friendly conversational workflows.
Finding this useful?
The AG2 team is working hard to create content like this, not to mention building a powerful, open-source, end-to-end platform for multi-agent automation.
The easiest way to show your support is just to star AG2 repo, but also take a look at it for contributions or simply to give it a try.
Also, let us know if you have any interesting use cases for ChatContext! Or maybe you would like to see more features or improvements? In any case, do join our Discord server for discussion.
Top comments (0)