Skip to content

Commit e28949a

Browse files
committed
Update agent system
1 parent 92b2368 commit e28949a

13 files changed

+131
-1045
lines changed

text_2_sql/autogen/agentic_text_2_sql.ipynb

Lines changed: 6 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
"cells": [
33
{
44
"cell_type": "code",
5-
"execution_count": 8,
5+
"execution_count": null,
66
"metadata": {},
77
"outputs": [],
88
"source": [
@@ -13,7 +13,7 @@
1313
},
1414
{
1515
"cell_type": "code",
16-
"execution_count": 9,
16+
"execution_count": null,
1717
"metadata": {},
1818
"outputs": [],
1919
"source": [
@@ -22,95 +22,18 @@
2222
},
2323
{
2424
"cell_type": "code",
25-
"execution_count": 10,
25+
"execution_count": null,
2626
"metadata": {},
27-
"outputs": [
28-
{
29-
"data": {
30-
"text/plain": [
31-
"True"
32-
]
33-
},
34-
"execution_count": 10,
35-
"metadata": {},
36-
"output_type": "execute_result"
37-
}
38-
],
27+
"outputs": [],
3928
"source": [
4029
"dotenv.load_dotenv()"
4130
]
4231
},
4332
{
4433
"cell_type": "code",
45-
"execution_count": 11,
34+
"execution_count": null,
4635
"metadata": {},
47-
"outputs": [
48-
{
49-
"name": "stderr",
50-
"output_type": "stream",
51-
"text": [
52-
"INFO:autogen_core:Publishing message of type GroupChatStart to all subscribers: {'message': TextMessage(source='user', models_usage=None, content='What are the total number of sales within 2008?')}\n",
53-
"INFO:autogen_core:Calling message handler for SQL_Query_Generation_Agent with message type GroupChatStart published by Unknown\n",
54-
"INFO:autogen_core:Calling message handler for SQL_Schema_Selection_Agent with message type GroupChatStart published by Unknown\n",
55-
"INFO:autogen_core:Calling message handler for SQL_Query_Correction_Agent with message type GroupChatStart published by Unknown\n",
56-
"INFO:autogen_core:Calling message handler for SQL_Query_Cache_Agent with message type GroupChatStart published by Unknown\n",
57-
"INFO:autogen_core:Calling message handler for Answer_Agent with message type GroupChatStart published by Unknown\n",
58-
"INFO:autogen_core:Calling message handler for Question_Decomposition_Agent with message type GroupChatStart published by Unknown\n",
59-
"INFO:autogen_core:Calling message handler for SQL_Query_Generation_Agent with message type GroupChatStart published by Unknown\n",
60-
"INFO:autogen_core:Calling message handler for SQL_Schema_Selection_Agent with message type GroupChatStart published by Unknown\n",
61-
"INFO:autogen_core:Calling message handler for SQL_Query_Correction_Agent with message type GroupChatStart published by Unknown\n",
62-
"INFO:autogen_core:Calling message handler for SQL_Query_Cache_Agent with message type GroupChatStart published by Unknown\n",
63-
"INFO:autogen_core:Calling message handler for Answer_Agent with message type GroupChatStart published by Unknown\n",
64-
"INFO:autogen_core:Calling message handler for Question_Decomposition_Agent with message type GroupChatStart published by Unknown\n",
65-
"INFO:autogen_core:Calling message handler for group_chat_manager with message type GroupChatStart published by Unknown\n",
66-
"INFO:autogen_core:Publishing message of type GroupChatStart to all subscribers: {'message': TextMessage(source='user', models_usage=None, content='What are the total number of sales within 2008?')}\n",
67-
"INFO:autogen_core:Calling message handler for collect_output_messages with message type GroupChatStart published by group_chat_manager/3b4d5ab6-7cfa-4179-8079-28bf0d1e0db4\n",
68-
"INFO:autogen_agentchat.events:source='user' models_usage=None content='What are the total number of sales within 2008?'\n",
69-
"INFO:httpx:HTTP Request: POST https://aoai-text2sql-adi.openai.azure.com/openai/deployments/gpt-4o-mini/chat/completions?api-version=2024-08-01-preview \"HTTP/1.1 200 OK\"\n",
70-
"INFO:autogen_core.events:{\"prompt_tokens\": 1801, \"completion_tokens\": 6, \"type\": \"LLMCall\"}\n",
71-
"INFO:autogen_core:Publishing message of type GroupChatRequestPublish to all subscribers: {}\n",
72-
"INFO:autogen_core:Calling message handler for SQL_Query_Cache_Agent with message type GroupChatRequestPublish published by group_chat_manager/3b4d5ab6-7cfa-4179-8079-28bf0d1e0db4\n",
73-
"INFO:httpx:HTTP Request: POST https://aoai-text2sql-adi.openai.azure.com/openai/deployments/gpt-4o-mini/chat/completions?api-version=2024-08-01-preview \"HTTP/1.1 200 OK\"\n",
74-
"INFO:autogen_core.events:{\"prompt_tokens\": 1582, \"completion_tokens\": 28, \"type\": \"LLMCall\"}\n",
75-
"INFO:autogen_core:Publishing message of type GroupChatMessage to all subscribers: {'message': ToolCallMessage(source='SQL_Query_Cache_Agent', models_usage=RequestUsage(prompt_tokens=1582, completion_tokens=28), content=[FunctionCall(id='call_CSvANg2842Hg57erUpUGssoe', arguments='{\"question\":\"What is the total number of sales orders in 2008?\"}', name='fetch_queries_from_cache')])}\n",
76-
"INFO:autogen_core:Calling message handler for collect_output_messages with message type GroupChatMessage published by SQL_Query_Cache_Agent/3b4d5ab6-7cfa-4179-8079-28bf0d1e0db4\n",
77-
"INFO:autogen_agentchat.events:source='SQL_Query_Cache_Agent' models_usage=RequestUsage(prompt_tokens=1582, completion_tokens=28) content=[FunctionCall(id='call_CSvANg2842Hg57erUpUGssoe', arguments='{\"question\":\"What is the total number of sales orders in 2008?\"}', name='fetch_queries_from_cache')]\n",
78-
"INFO:httpx:HTTP Request: POST https://aoai-text2sql-adi.openai.azure.com/openai/deployments/text-embedding-ada-002/embeddings?api-version=2023-03-15-preview \"HTTP/1.1 200 OK\"\n",
79-
"INFO:azure.core.pipeline.policies.http_logging_policy:Request URL: 'https://aisearch-text2sql-adi.search.windows.net/indexes('text-2-sql-query-cache-index')/docs/search.post.search?api-version=REDACTED'\n",
80-
"Request method: 'POST'\n",
81-
"Request headers:\n",
82-
" 'Content-Type': 'application/json'\n",
83-
" 'Content-Length': '34704'\n",
84-
" 'api-key': 'REDACTED'\n",
85-
" 'Accept': 'application/json;odata.metadata=none'\n",
86-
" 'x-ms-client-request-id': '6a968312-a735-11ef-a762-0242ac110002'\n",
87-
" 'User-Agent': 'azsdk-python-search-documents/11.6.0b5 Python/3.12.6 (Linux-5.15.153.1-microsoft-standard-WSL2-x86_64-with-glibc2.36)'\n",
88-
"A body is sent with the request\n",
89-
"INFO:azure.core.pipeline.policies.http_logging_policy:Response status: 200\n",
90-
"Response headers:\n",
91-
" 'Transfer-Encoding': 'chunked'\n",
92-
" 'Content-Type': 'application/json; odata.metadata=none; odata.streaming=true; charset=utf-8'\n",
93-
" 'Content-Encoding': 'REDACTED'\n",
94-
" 'Vary': 'REDACTED'\n",
95-
" 'Server': 'Microsoft-IIS/10.0'\n",
96-
" 'Strict-Transport-Security': 'REDACTED'\n",
97-
" 'Preference-Applied': 'REDACTED'\n",
98-
" 'OData-Version': 'REDACTED'\n",
99-
" 'request-id': '6a968312-a735-11ef-a762-0242ac110002'\n",
100-
" 'elapsed-time': 'REDACTED'\n",
101-
" 'Strict-Transport-Security': 'REDACTED'\n",
102-
" 'Date': 'Wed, 20 Nov 2024 11:48:53 GMT'\n",
103-
"INFO:root:Results: []\n",
104-
"INFO:autogen_core:Publishing message of type GroupChatMessage to all subscribers: {'message': ToolCallResultMessage(source='SQL_Query_Cache_Agent', models_usage=None, content=[FunctionExecutionResult(content=\"{'cached_questions': None}\", call_id='call_CSvANg2842Hg57erUpUGssoe')])}\n",
105-
"INFO:autogen_core:Calling message handler for collect_output_messages with message type GroupChatMessage published by SQL_Query_Cache_Agent/3b4d5ab6-7cfa-4179-8079-28bf0d1e0db4\n",
106-
"INFO:autogen_agentchat.events:source='SQL_Query_Cache_Agent' models_usage=None content=[FunctionExecutionResult(content=\"{'cached_questions': None}\", call_id='call_CSvANg2842Hg57erUpUGssoe')]\n",
107-
"INFO:httpx:HTTP Request: POST https://aoai-text2sql-adi.openai.azure.com/openai/deployments/gpt-4o-mini/chat/completions?api-version=2024-08-01-preview \"HTTP/1.1 500 model_error\"\n",
108-
"INFO:openai._base_client:Retrying request to /chat/completions in 0.951392 seconds\n",
109-
"INFO:httpx:HTTP Request: POST https://aoai-text2sql-adi.openai.azure.com/openai/deployments/gpt-4o-mini/chat/completions?api-version=2024-08-01-preview \"HTTP/1.1 500 model_error\"\n",
110-
"INFO:openai._base_client:Retrying request to /chat/completions in 1.995870 seconds\n"
111-
]
112-
}
113-
],
36+
"outputs": [],
11437
"source": [
11538
"result = await SELECTOR.run(task=\"What are the total number of sales within 2008?\")"
11639
]
@@ -124,13 +47,6 @@
12447
"print(result.messages[-1].content)"
12548
]
12649
},
127-
{
128-
"cell_type": "code",
129-
"execution_count": null,
130-
"metadata": {},
131-
"outputs": [],
132-
"source": []
133-
},
13450
{
13551
"cell_type": "code",
13652
"execution_count": null,
Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,20 @@
11
model:
22
gpt-4o-mini
33
description:
4-
An agent that takes the final results from the SQL query and writes the answer to the user's question
4+
"An agent that takes the final results from the SQL query and writes the answer to the user's question"
55
system_message:
6-
Write a data-driven answer that directly addresses the user's question. Use the results from the SQL query to provide the answer. Do not make up or guess the answer.
6+
"Write a data-driven answer that directly addresses the user's question. Use the results from the SQL query to provide the answer. Do not make up or guess the answer.
77
8-
End your answer with 'TERMINATE'
8+
Return your answer in the following format:
9+
10+
{
11+
'answer': '<GENERATED ANSWER>',
12+
'sources': [
13+
{'title': <SOURCE SCHEMA NAME 1>, 'chunk': <SOURCE 1 CONTEXT CHUNK>, 'reference': '<SOURCE 1 SQL QUERY>'},
14+
{'title': <SOURCE SCHEMA NAME 2>, 'chunk': <SOURCE 2 CONTEXT CHUNK>, 'reference': '<SOURCE 2 SQL QUERY>'}
15+
]
16+
}
17+
18+
Title is the entity name of the schema, chunk is the result of the SQL query and reference is the SQL query used to generate the answer.
19+
20+
End your answer with 'TERMINATE'"
Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
model:
22
gpt-4o-mini
33
description:
4-
An agent that will decompose the user's question into smaller parts to be used in the SQL queries. Use this agent when the user's question is too complex to be answered in one SQL query. Only use if the user's question is too complex to be answered in one SQL query.
4+
"An agent that will decompose the user's question into smaller parts to be used in the SQL queries. Use this agent when the user's question is too complex to be answered in one SQL query. Only use if the user's question is too complex to be answered in one SQL query.
55
6-
Only use this agent once per user question and after the 'Query Cache Agent' if the results are none.
6+
Only use this agent once per user question and after the 'Query Cache Agent' if the results are none."
77
system_message:
8-
You are a helpful AI Assistant that specialises in decomposing complex user questions into smaller parts that can be used in SQL queries.
8+
"You are a helpful AI Assistant that specialises in decomposing complex user questions into smaller parts that can be used in SQL queries.
99
10-
Break down the user's question into smaller parts that can be used in SQL queries.
10+
Break down the user's question into smaller parts that can be used in SQL queries."
Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
model:
22
gpt-4o-mini
33
description:
4-
An agent that can take a user's question and use a pre-built SQL query cache to quickly fetch results for the user. Use this agent before all other agents to quickly fetch results for the user. If the results are not as expected, use the SQL Query Correction Agent to correct the SQL query if the schemas are similar. If no results are found, use the SQL Schema Selection Agent to select the correct schema.
4+
"An agent that can take a user's question and use a pre-built SQL query cache to quickly fetch results for the user. Use this agent before all other agents to quickly fetch results for the user. If the results are not as expected, use the SQL Query Correction Agent to correct the SQL query if the schemas are similar. If no results are found, use the SQL Schema Selection Agent to select the correct schema.
55
6-
Only use this agent once per user question and at the beginning of the conversation.
6+
Only use this agent once per user question and at the beginning of the conversation."
77
system_message:
8-
You are a helpful AI Assistant that specialises in using a pre-built SQL query cache to quickly fetch results for the user.
8+
"You are a helpful AI Assistant that specialises in using a pre-built SQL query cache to quickly fetch results for the user.
99
10-
Use the tool and access to the pre-built query cache to ask the pass the user's question to the cache and fetch the results.
10+
Use the tool and access to the pre-built query cache to ask the pass the user's question to the cache and fetch the results."
1111
tools:
1212
- sql_query_cache_tool
Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,19 @@
11
model:
22
gpt-4o-mini
33
description:
4-
An agent that will look at the SQL query, SQL query results and correct any mistakes in the SQL query to ensure the correct results are returned. Use this agent AFTER the SQL query has been executed and the results are not as expected.
4+
"An agent that will look at the SQL query, SQL query results and correct any mistakes in the SQL query to ensure the correct results are returned. Use this agent AFTER the SQL query has been executed and the results are not as expected."
55
system_message:
6-
You are a helpful AI Assistant that specialises in correcting invalid SQL queries or queries that do not return the expected results.
6+
"You are a helpful AI Assistant that specialises in correcting invalid SQL queries or queries that do not return the expected results.
77
8-
Review the SQL query provided and correct any errors or issues that you find. Ensure that the corrected query returns the expected results.
8+
Review the SQL query provided and correct any errors or issues that you find. Bear in mind that the target database engine is {{ target_engine }}, SQL queries must be able compatible to run on {{ target_engine }} {{ engine_specific_rules }}
99
10-
If you are consistently unable to correct the SQL query and cannot use the schemas to answer the question. Say 'I am unable to correct the SQL query. Please ask another question.' and then end your answer with 'TERMINATE'
10+
Ensure that the corrected query returns the expected results in context of the question.
11+
12+
If there are no errors and the SQL query is correct, return 'VALIDATED'.
13+
14+
If the SQL query needs adjustment, correct the SQL query and provide the corrected SQL query and then run the query.
15+
16+
If you are consistently unable to correct the SQL query and cannot use the schemas to answer the question. Say 'I am unable to correct the SQL query. Please ask another question.' and then end your answer with 'TERMINATE'"
1117
tools:
1218
- sql_get_entity_schemas_tool
1319
- sql_query_execution_tool
Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
model:
22
gpt-4o-mini
33
description:
4-
An agent that can generate SQL queries once given the schema and the user's question. It will run the SQL query to fetch the results. Use this agent after the SQL Schema Selection Agent has selected the correct schema.
4+
"An agent that can generate SQL queries once given the schema and the user's question. It will run the SQL query to fetch the results. Use this agent after the SQL Schema Selection Agent has selected the correct schema."
55
system_message:
66
"You are a helpful AI Assistant that specialises in writing and executing SQL Queries to answer a given user's question.
77
@@ -10,9 +10,9 @@ system_message:
1010
Only use schema / column information provided when constructing a SQL query. Do not use any other entities and columns in your SQL query, other than those defined above.
1111
Do not makeup or guess column names.
1212
13-
The target database engine is SQL Server, SQL queries must be able compatible to run on Target Engine.
13+
The target database engine is {{ target_engine }}, SQL queries must be able compatible to run on {{ target_engine }} {{ engine_specific_rules }}
1414
You must only provide SELECT SQL queries.
15-
For a given entity, use the 'SelectFromEntity' property returned in the schema in the SELECT FROM part of the SQL query. If the property is {{'SelectFromEntity': 'test_schema.test_table'}}, the select statement will be formulated from 'SELECT <VALUES> FROM test_schema.test_table WHERE <CONDITION>.
15+
For a given entity, use the 'SelectFromEntity' property returned in the schema in the SELECT FROM part of the SQL query. If the property is {'SelectFromEntity': 'test_schema.test_table'}, the select statement will be formulated from 'SELECT <VALUES> FROM test_schema.test_table WHERE <CONDITION>.
1616
1717
If you don't know how the value is formatted in a column, run a query against the column to get the unique values that might match your query or use the corresponding lookup values. Use a 'like' operator to match the values, rather than a direct match unless you are sure of the value.
1818
Some columns in the schema may have the properties 'AllowedValues' or 'SampleValues'. Use these values to determine the possible values that can be used in the SQL query.
@@ -22,3 +22,4 @@ system_message:
2222
Always run any SQL query you generate to return the results."
2323
tools:
2424
- sql_query_execution_tool
25+
- sql_get_entity_schemas_tool
Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
model:
22
gpt-4o-mini
33
description:
4-
An agent that can take a user's question and extract the schema of a view or table in the SQL Database by selecting the most relevant entity based on the search term.
4+
"An agent that can take a user's question and extract the schema of a view or table in the SQL Database by selecting the most relevant entity based on the search term.
55
6-
Call this in parallel if needed multiple times. Limit the use of this agent where possible.
6+
Call this in parallel if needed multiple times. Limit the use of this agent where possible."
77
system_message:
8-
You are a helpful AI Assistant that specialises in selecting relevant SQL schemas to answer a given user's question.
8+
"You are a helpful AI Assistant that specialises in selecting relevant SQL schemas to answer a given user's question.
99
10-
Use the tools available to you to select the correct schemas that will help. Extract key terms from the user's question and use them to search for the correct schema.
10+
Use the tools available to you to select the correct schemas that will help. Extract key terms from the user's question and use them to search for the correct schema."
1111
tools:
1212
- sql_get_entity_schemas_tool

text_2_sql/autogen/requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,3 +7,4 @@ azure-search-documents==11.6.0b5
77
azure-identity
88
python-dotenv
99
openai
10+
jinja2

0 commit comments

Comments
 (0)