I've been testing AutoGen for automating some of our weekly marketing performance reporting. The core challenge was getting the agents to query our internal Postgres database, which houses our unified marketing funnel (web sessions, leads, campaign touches, revenue). The documentation on custom connections is a bit scattered, so here's a consolidated, operational guide based on what worked.
The goal was to have a `reporting_agent` that could execute SQL against specific schemas and return results to a `user_proxy` for formatting. The key is in the `llm_config` and defining a custom function for database access.
**Step 1: Define the Database Connection Function**
You'll need `psycopg2` (or `asyncpg`) installed. This function will be passed to the agent.
```python
import psycopg2
from psycopg2 import sql
def execute_postgres_query(query: str):
"""
Executes a read-only query against the marketing_db.
Returns results as a list of tuples or an error message.
"""
conn = None
try:
conn = psycopg2.connect(
host="your.internal.host",
database="marketing_db",
user="readonly_user",
password="your_password",
port=5432
)
cur = conn.cursor()
cur.execute(query)
if query.strip().lower().startswith("select"):
results = cur.fetchall()
col_names = [desc[0] for desc in cur.description]
return {"columns": col_names, "data": results}
else:
conn.commit()
return {"message": "Query executed successfully."}
except Exception as e:
return {"error": str(e)}
finally:
if conn:
conn.close()
```
**Step 2: Configure the Agent with Function Mapping**
Crucially, you must add the function to the `function_map` in `llm_config` and also list it in `functions`.
```python
from autogen import ConversableAgent
reporting_agent = ConversableAgent(
name="reporting_agent",
system_message="You are a SQL expert. You generate and execute PostgreSQL queries to answer data questions. You only use the `execute_postgres_query` function. You return the results succinctly.",
llm_config={
"config_list": [{"model": "gpt-4", "api_key": "your_key"}],
"functions": [
{
"name": "execute_postgres_query",
"description": "Execute a SQL query against the marketing database and return results.",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The SQL query to execute.",
}
},
"required": ["query"],
},
}
],
"function_map": {
"execute_postgres_query": execute_postgres_query
}
}
)
```
**Step 3: Implement Safety and Context**
To prevent broad `SELECT *` and confine queries to relevant schemas, I added a system prompt directive and a pre-execution validation layer within the `execute_postgres_query` function. For example:
- Check for `DROP`, `DELETE`, `INSERT`, `UPDATE` and reject.
- Prepend a default schema (e.g., `SET search_path TO 'marketing';`) or validate table names against a whitelist.
**Pitfalls Encountered:**
* The agent sometimes tries to invent tables. Providing a clear data dictionary in the system prompt helps.
* Connection pooling isn't handled here; for frequent use, consider a connection pooler like `pgbouncer` or a dedicated middleware API.
* Large result sets can timeout. Implementing `LIMIT` logic in the function or agent instructions is advisable.
This setup now automatically generates weekly top-of-funnel and lead source performance summaries. The next step is to integrate this with a `chart_agent` that can take the result sets and generate visualization code.