This commit is contained in:
isaac hershenson
2024-06-18 14:45:06 -07:00
parent 004be5a644
commit ec8fca977e
17 changed files with 202 additions and 125 deletions
+54 -26
View File
@@ -43,11 +43,11 @@
"execution_count": 1,
"id": "6c05a600f1afb5b6",
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T21:24:00.532147Z",
"start_time": "2024-06-12T21:24:00.526043Z"
}
},
"collapsed": false
},
"outputs": [],
"source": [
@@ -78,11 +78,11 @@
"execution_count": 2,
"id": "64b0bf1b14c2e902",
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T21:24:09.918436Z",
"start_time": "2024-06-12T21:24:09.608563Z"
}
},
"collapsed": false
},
"outputs": [
{
@@ -125,11 +125,11 @@
"execution_count": 3,
"id": "a60191bd3489f278",
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T21:24:14.663745Z",
"start_time": "2024-06-12T21:24:13.527958Z"
}
},
"collapsed": false
},
"outputs": [],
"source": [
@@ -142,11 +142,11 @@
"execution_count": 4,
"id": "1f1e1f4f86ed54",
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T21:24:15.891582Z",
"start_time": "2024-06-12T21:24:15.289782Z"
}
},
"collapsed": false
},
"outputs": [
{
@@ -192,11 +192,11 @@
"execution_count": 5,
"id": "deae8460e4cf72b1",
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T21:24:17.557848Z",
"start_time": "2024-06-12T21:24:17.508550Z"
}
},
"collapsed": false
},
"outputs": [],
"source": [
@@ -214,6 +214,7 @@
" [RunnableLambda(handle_tool_error)], exception_key=\"error\"\n",
" )\n",
"\n",
"\n",
"def handle_tool_error(state) -> dict:\n",
" error = state.get(\"error\")\n",
" tool_calls = state[\"messages\"][-1].tool_calls\n",
@@ -341,6 +342,7 @@
" return \"Error: Query failed. Please rewrite your query and try again.\"\n",
" return result\n",
"\n",
"\n",
"print(db_query_tool.invoke(\"SELECT * FROM Artist LIMIT 10;\"))"
]
},
@@ -393,8 +395,12 @@
"\n",
"You will call the appropriate tool to execute the query after running this check.\"\"\"\n",
"\n",
"query_check_prompt = ChatPromptTemplate.from_messages([(\"system\", query_check_system),(\"placeholder\", \"{messages}\")])\n",
"query_check = query_check_prompt | ChatOpenAI(model=\"gpt-4o\", temperature=0).bind_tools([db_query_tool], tool_choice=\"required\")\n",
"query_check_prompt = ChatPromptTemplate.from_messages(\n",
" [(\"system\", query_check_system), (\"placeholder\", \"{messages}\")]\n",
")\n",
"query_check = query_check_prompt | ChatOpenAI(model=\"gpt-4o\", temperature=0).bind_tools(\n",
" [db_query_tool], tool_choice=\"required\"\n",
")\n",
"\n",
"query_check.invoke({\"messages\": [(\"user\", \"SELET * FROM Artist LIMIT 10;\")]})"
]
@@ -440,9 +446,11 @@
"class State(TypedDict):\n",
" messages: Annotated[list[AnyMessage], add_messages]\n",
"\n",
"\n",
"# Define a new graph\n",
"workflow = StateGraph(State)\n",
"\n",
"\n",
"# Add a node for the first tool call\n",
"def first_tool_call(state: State) -> dict[str, list[AIMessage]]:\n",
" return {\n",
@@ -452,8 +460,7 @@
" tool_calls=[\n",
" {\n",
" \"name\": \"sql_db_list_tables\",\n",
" \"args\": {\n",
" },\n",
" \"args\": {},\n",
" \"id\": \"tool_abcd123\",\n",
" }\n",
" ],\n",
@@ -461,31 +468,41 @@
" ]\n",
" }\n",
"\n",
"\n",
"def model_check_query(state: State) -> dict[str, list[AIMessage]]:\n",
" \"\"\"\n",
" Use this tool to double-check if your query is correct before executing it.\n",
" \"\"\"\n",
" return {\n",
" \"messages\": [\n",
" query_check.invoke({\"messages\": [state[\"messages\"][-1]]})\n",
" ]\n",
" }\n",
" return {\"messages\": [query_check.invoke({\"messages\": [state[\"messages\"][-1]]})]}\n",
"\n",
"\n",
"workflow.add_node(\"first_tool_call\", first_tool_call)\n",
"\n",
"# Add nodes for the first two tools\n",
"workflow.add_node(\"list_tables_tool\", create_tool_node_with_fallback([list_tables_tool]))\n",
"workflow.add_node(\n",
" \"list_tables_tool\", create_tool_node_with_fallback([list_tables_tool])\n",
")\n",
"workflow.add_node(\"get_schema_tool\", create_tool_node_with_fallback([get_schema_tool]))\n",
"\n",
"# Add a node for a model to choose the relevant tables based on the question and available tables\n",
"model_get_schema = ChatOpenAI(model=\"gpt-4o\", temperature=0).bind_tools([get_schema_tool])\n",
"workflow.add_node(\"model_get_schema\", lambda state: {\"messages\": [model_get_schema.invoke(state[\"messages\"])],})\n",
"model_get_schema = ChatOpenAI(model=\"gpt-4o\", temperature=0).bind_tools(\n",
" [get_schema_tool]\n",
")\n",
"workflow.add_node(\n",
" \"model_get_schema\",\n",
" lambda state: {\n",
" \"messages\": [model_get_schema.invoke(state[\"messages\"])],\n",
" },\n",
")\n",
"\n",
"\n",
"# Describe a tool to represent the end state\n",
"class SubmitFinalAnswer(BaseModel):\n",
" \"\"\"Submit the final answer to the user based on the query results.\"\"\"\n",
"\n",
" final_answer: str = Field(..., description=\"The final answer to the user\")\n",
"\n",
"\n",
"# Add a node for a model to generate a query based on the question and schema\n",
"query_gen_system = \"\"\"You are a SQL expert with a strong attention to detail.\n",
"\n",
@@ -509,11 +526,17 @@
"If you have enough information to answer the input question, simply invoke the appropriate tool to submit the final answer to the user.\n",
"\n",
"DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.\"\"\"\n",
"query_gen_prompt = ChatPromptTemplate.from_messages([(\"system\", query_gen_system),(\"placeholder\", \"{messages}\")])\n",
"query_gen = query_gen_prompt | ChatOpenAI(model=\"gpt-4o\", temperature=0).bind_tools([SubmitFinalAnswer])\n",
"query_gen_prompt = ChatPromptTemplate.from_messages(\n",
" [(\"system\", query_gen_system), (\"placeholder\", \"{messages}\")]\n",
")\n",
"query_gen = query_gen_prompt | ChatOpenAI(model=\"gpt-4o\", temperature=0).bind_tools(\n",
" [SubmitFinalAnswer]\n",
")\n",
"\n",
"\n",
"def query_gen_node(state: State):\n",
" message = query_gen.invoke(state)\n",
" \n",
"\n",
" # Sometimes, the LLM will hallucinate and call the wrong tool. We need to catch this and return an error message.\n",
" tool_messages = []\n",
" if message.tool_calls:\n",
@@ -529,6 +552,7 @@
" tool_messages = []\n",
" return {\"messages\": [message] + tool_messages}\n",
"\n",
"\n",
"workflow.add_node(\"query_gen\", query_gen_node)\n",
"\n",
"# Add a node for the model to check the query before executing it\n",
@@ -537,6 +561,7 @@
"# Add node for executing the query\n",
"workflow.add_node(\"execute_query\", create_tool_node_with_fallback([db_query_tool]))\n",
"\n",
"\n",
"# Define a conditional edge to decide whether to continue or end the workflow\n",
"def should_continue(state: State) -> Literal[END, \"correct_query\", \"query_gen\"]:\n",
" messages = state[\"messages\"]\n",
@@ -549,6 +574,7 @@
" else:\n",
" return \"correct_query\"\n",
"\n",
"\n",
"# Specify the edges between the nodes\n",
"workflow.set_entry_point(\"first_tool_call\")\n",
"workflow.add_edge(\"first_tool_call\", \"list_tables_tool\")\n",
@@ -649,7 +675,9 @@
}
],
"source": [
"for event in app.stream({\"messages\": [(\"user\", \"Which sales agent made the most in sales in 2009?\")]}):\n",
"for event in app.stream(\n",
" {\"messages\": [(\"user\", \"Which sales agent made the most in sales in 2009?\")]}\n",
"):\n",
" print(event)"
]
}