Tools with Dependency Injection

Tools Dependency Injection

Dependency Injection is a secure way to connect external functions to agents without exposing sensitive data such as passwords, tokens, or personal information. This approach ensures that sensitive information remains protected while still allowing agents to perform their tasks effectively, even when working with large language models (LLMs).

In this guide, we’ll explore how to build secure workflows that handle sensitive data safely.

As an example, we’ll create an agent that retrieves user’s account balance. The best part is that sensitive data like username and password are never shared with the LLM. Instead, it’s securely injected directly into the function at runtime, keeping it safe while maintaining seamless functionality.

Why Dependency Injection Is Essential

Here’s why dependency injection is a game-changer for secure LLM workflows:

• Enhanced Security: Your sensitive data is never directly exposed to the LLM.
• Simplified Development: Secure data can be seamlessly accessed by functions without requiring complex configurations.
• Unmatched Flexibility: It supports safe integration of diverse workflows, allowing you to scale and adapt with ease.

In this guide, we’ll explore how to set up dependency injection and build secure workflows. Let’s dive in!

Installation

To install AG2, simply run the following command:

pip install ag2

Imports

The functionality demonstrated in this guide is located in the autogen.tools.dependency_injection module. This module provides key components for dependency injection:

• BaseContext: abstract base class used to define and encapsulate data contexts, such as user account information, which can then be injected into functions or agents securely.
• Depends: a function used to declare and inject dependencies, either from a context (like BaseContext) or a function, ensuring sensitive data is provided securely without direct exposure.

import os
from typing import Annotated, Literal

from pydantic import BaseModel

from autogen import GroupChat, GroupChatManager
from autogen.agentchat import ConversableAgent, UserProxyAgent
from autogen.tools.dependency_injection import BaseContext, Depends

Define a BaseContext Class

We start by defining a BaseContext class for accounts. This will act as the base structure for dependency injection. By using this approach, sensitive information like usernames and passwords is never exposed to the LLM.

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

To ensure that the provided account is valid and retrieve its balance, we create two helper functions.

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 to NEVER).
• AssistantAgent represents the AI agent, configured with the LLM settings.

config_list = [{"model": "gpt-4o-mini", "api_key": os.environ["OPENAI_API_KEY"]}]

assistant = ConversableAgent(
    name="assistant",
    llm_config={"config_list": config_list},
)
user_proxy = UserProxyAgent(
    name="user_proxy_1",
    human_input_mode="NEVER",
    llm_config=False,
)

Injecting a BaseContext Parameter

In the example below we register the function and use dependency injection to automatically inject the bob_account Account object into the function. This account parameter will not be visible to the LLM.

Note: You can also use account: Account = Depends(bob_account) as an alternative syntax.

@user_proxy.register_for_execution()
@assistant.register_for_llm(description="Get the balance of the account")
def get_balance_1(
    # Account which will be injected to the function
    account: Annotated[Account, Depends(bob_account)],
    # It is also possible to use the following syntax to define the dependency
    # account: Account = Depends(bob_account),
) -> str:
    return _get_balance(account)

Finally, we initiate a chat to retrieve the balance.

user_proxy.initiate_chat(assistant, message="Get the user's account balance", max_turns=2)

Injecting Parameters Without BaseContext

Sometimes, you might not want to use BaseContext. Here’s how to inject simple parameters directly.

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 to NEVER).
• AssistantAgent represents the AI agent, configured with the LLM settings.

config_list = [{"model": "gpt-4o-mini", "api_key": os.environ["OPENAI_API_KEY"]}]
assistant = ConversableAgent(
    name="assistant",
    llm_config={"config_list": config_list},
)
user_proxy = UserProxyAgent(
    name="user_proxy_1",
    human_input_mode="NEVER",
    llm_config=False,
)

Register the Function with Direct Parameter Injection

Instead of injecting a full context like Account, you can directly inject individual parameters, such as the username and password, into a function. This allows for more granular control over the data injected into the function, and still ensures that sensitive information is managed securely.

Here’s how you can set it up:

def get_username() -> str:
    return "bob"


def get_password() -> str:
    return "password456"


@user_proxy.register_for_execution()
@assistant.register_for_llm(description="Get the balance of the account")
def get_balance_2(
    username: Annotated[str, Depends(get_username)],
    password: Annotated[str, Depends(get_password)],
    # or use lambdas
    # username: Annotated[str, Depends(lambda: "bob")],
    # password: Annotated[str, Depends(lambda: "password456")],
) -> str:
    account = Account(username=username, password=password)
    return _get_balance(account)

Initiate the Chat

As before, initiate a chat to test the function.

user_proxy.initiate_chat(assistant, message="Get users balance", max_turns=2)

Aligning Contexts to Agents

You can match specific dependencies, such as 3rd party system credentials, with specific agents by using tools with dependency injection.

In this example we have 2 external systems and have 2 related login credentials. We don’t want or need the LLM to be aware of these credentials.

Mock third party systems

# Mock third party system functions
# Imagine that these use the username and password to authenticate


def weather_api_call(username: str, password: str, location: str) -> str:
    print(f"Accessing third party Weather System using username {username}")
    return "It's sunny and 40 degrees Celsius in Sydney, Australia."


def my_ticketing_system_availability(username: str, password: str, concert: str) -> bool:
    print(f"Accessing third party Ticketing System using username {username}")
    return False

Create Agents

weather_agent = ConversableAgent(
    name="weather_agent",
    system_message="You are a Weather Agent, you can only get the weather.",
    description="Weather Agent solely used for getting weather.",
    llm_config={"config_list": config_list},
)

ticket_agent = ConversableAgent(
    name="ticket_agent",
    system_message="You are a Ticketing Agent, you can only get ticket availability.",
    description="Ticketing Agent solely used for getting ticket availability.",
    llm_config={"config_list": config_list},
)
user_proxy = UserProxyAgent(
    name="user_proxy_1",
    human_input_mode="NEVER",
    llm_config=False,
)

Create BaseContext class for credentials

# We create a class based on BaseContext
class ThirdPartyCredentials(BaseContext, BaseModel):
    username: str
    password: str

Create Credentials and Functions with Dependency Injection

# Weather API
weather_account = ThirdPartyCredentials(username="ag2weather", password="wbkvEehV1A")


@user_proxy.register_for_execution()
@weather_agent.register_for_llm(description="Get the weather for a location")
def get_weather(
    location: str,
    credentials: Annotated[ThirdPartyCredentials, Depends(weather_account)],
) -> str:
    # Access the Weather API using the credentials
    return weather_api_call(username=credentials.username, password=credentials.password, location=location)


# Ticketing System
ticket_system_account = ThirdPartyCredentials(username="ag2tickets", password="EZRIVeVWvA")


@user_proxy.register_for_execution()
@ticket_agent.register_for_llm(description="Get the availability of tickets for a concert")
def tickets_available(
    concert_name: str,
    credentials: Annotated[ThirdPartyCredentials, Depends(ticket_system_account)],
) -> bool:
    return my_ticketing_system_availability(
        username=credentials.username, password=credentials.password, concert=concert_name
    )

Create Group Chat and run

groupchat = GroupChat(agents=[user_proxy, weather_agent, ticket_agent], messages=[], max_round=5)
manager = GroupChatManager(groupchat=groupchat, llm_config={"config_list": config_list})

message = (
    "Start by getting the weather for Sydney, Australia, and follow that up by checking "
    "if there are tickets for the 'AG2 Live' concert."
)
user_proxy.initiate_chat(manager, message=message, max_turns=1)