Files
docs/build/snippets/python/code-samples/langgraph-sql-agent-define-steps-js.mdx
T
2026-07-29 10:28:19 +00:00

129 lines
3.9 KiB
Plaintext

```ts
import {
AIMessage,
HumanMessage,
SystemMessage,
ToolMessage,
} from "@langchain/core/messages";
import { ToolNode } from "@langchain/langgraph/prebuilt";
import {
END,
GraphNode,
MessagesValue,
START,
StateGraph,
StateSchema,
} from "@langchain/langgraph";
// Create tool nodes for schema and query execution
const getSchemaNode = new ToolNode([getSchemaTool]);
const runQueryNode = new ToolNode([queryTool]);
// Define state schema
const MessagesState = new StateSchema({
messages: MessagesValue,
});
// Example: create a predetermined tool call
const listTables: GraphNode<typeof MessagesState> = async (state) => {
const toolCall = {
name: "sql_db_list_tables",
args: {},
id: "abc123",
type: "tool_call" as const,
};
const toolCallMessage = new AIMessage({
content: "",
tool_calls: [toolCall],
});
const toolMessage = await listTablesTool.invoke({});
const response = new AIMessage(`Available tables: ${toolMessage}`);
return {
messages: [
toolCallMessage,
new ToolMessage({ content: toolMessage, tool_call_id: "abc123" }),
response,
],
};
};
// Example: force a model to create a tool call
const callGetSchema: GraphNode<typeof MessagesState> = async (state) => {
const llmWithTools = model!.bindTools([getSchemaTool], {
tool_choice: "any",
});
const response = await llmWithTools.invoke(state.messages);
return { messages: [response] };
};
const topK = 5;
const generateQuerySystemPrompt = `
You are an agent designed to interact with a SQL database.
Given an input question, create a syntactically correct ${dialect}
query to run, then look at the results of the query and return the answer. Unless
the user specifies a specific number of examples they wish to obtain, always limit
your query to at most ${topK} results.
You can order the results by a relevant column to return the most interesting
examples in the database. Never query for all the columns from a specific table,
only ask for the relevant columns given the question.
DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.
`;
const generateQuery: GraphNode<typeof MessagesState> = async (state) => {
const systemMessage = new SystemMessage(generateQuerySystemPrompt);
// We do not force a tool call here, to allow the model to
// respond naturally when it obtains the solution.
const llmWithTools = model!.bindTools([queryTool]);
const response = await llmWithTools.invoke([
systemMessage,
...state.messages,
]);
return { messages: [response] };
};
const checkQuerySystemPrompt = `
You are a SQL expert with a strong attention to detail.
Double check the ${dialect} query for common mistakes, including:
- Using NOT IN with NULL values
- Using UNION when UNION ALL should have been used
- Using BETWEEN for exclusive ranges
- Data type mismatch in predicates
- Properly quoting identifiers
- Using the correct number of arguments for functions
- Casting to the correct data type
- Using the proper columns for joins
If there are any of the above mistakes, rewrite the query. If there are no mistakes,
just reproduce the original query.
You will call the appropriate tool to execute the query after running this check.
`;
const checkQuery: GraphNode<typeof MessagesState> = async (state) => {
const systemMessage = new SystemMessage(checkQuerySystemPrompt);
// Generate an artificial user message to check
const lastMessage = state.messages[state.messages.length - 1];
if (!lastMessage.tool_calls || lastMessage.tool_calls.length === 0) {
throw new Error("No tool calls found in the last message");
}
const toolCall = lastMessage.tool_calls[0];
const userMessage = new HumanMessage(toolCall.args.query);
const llmWithTools = model!.bindTools([queryTool], {
tool_choice: "any",
});
const response = await llmWithTools.invoke([systemMessage, userMessage]);
// Preserve the original message ID
response.id = lastMessage.id;
return { messages: [response] };
};
```