from langchain_community.agent_toolkits import SQLDatabaseToolkittoolkit = SQLDatabaseToolkit(db=db, llm=model)tools = toolkit.get_tools()for tool in tools: print(f"{tool.name}: {tool.description}\n")
sql_db_query: Input to this tool is a detailed and correct SQL query, output is a result from the database. If the query is not correct, an error message will be returned. If an error is returned, rewrite the query, check the query, and try again. If you encounter an issue with Unknown column 'xxxx' in 'field list', use sql_db_schema to query the correct table fields.sql_db_schema: Input to this tool is a comma-separated list of tables, output is the schema and sample rows for those tables. Be sure that the tables actually exist by calling sql_db_list_tables first! Example Input: table1, table2, table3sql_db_list_tables: Input is an empty string, output is a comma-separated list of tables in the database.sql_db_query_checker: Use this tool to double check if your query is correct before executing it. Always use this tool before executing a query with sql_db_query!
system_prompt = """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 userspecifies a specific number of examples they wish to obtain, always limit yourquery to at most {top_k} results.You can order the results by a relevant column to return the most interestingexamples in the database. Never query for all the columns from a specific table,only ask for the relevant columns given the question.You MUST double check your query before executing it. If you get an error whileexecuting a query, rewrite the query and try again.DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to thedatabase.To start you should ALWAYS look at the tables in the database to see what youcan query. Do NOT skip this step.Then you should query the schema of the most relevant tables.""".format( dialect=db.dialect, top_k=5,)
现在,使用模型、工具和提示创建代理:
from langchain.agents import create_agentagent = create_agent( model, tools, system_prompt=system_prompt,)
question = "Which genre on average has the longest tracks?"for step in agent.stream( {"messages": [{"role": "user", "content": question}]}, stream_mode="values",): step["messages"][-1].pretty_print()
================================ Human Message =================================Which genre on average has the longest tracks?================================== Ai Message ==================================Tool Calls: sql_db_list_tables (call_BQsWg8P65apHc8BTJ1NPDvnM) Call ID: call_BQsWg8P65apHc8BTJ1NPDvnM Args:================================= Tool Message =================================Name: sql_db_list_tablesAlbum, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track================================== Ai Message ==================================Tool Calls: sql_db_schema (call_i89tjKECFSeERbuACYm4w0cU) Call ID: call_i89tjKECFSeERbuACYm4w0cU Args: table_names: Track, Genre================================= Tool Message =================================Name: sql_db_schemaCREATE TABLE "Genre" ( "GenreId" INTEGER NOT NULL, "Name" NVARCHAR(120), PRIMARY KEY ("GenreId"))/*3 rows from Genre table:GenreId Name1 Rock2 Jazz3 Metal*/CREATE TABLE "Track" ( "TrackId" INTEGER NOT NULL, "Name" NVARCHAR(200) NOT NULL, "AlbumId" INTEGER, "MediaTypeId" INTEGER NOT NULL, "GenreId" INTEGER, "Composer" NVARCHAR(220), "Milliseconds" INTEGER NOT NULL, "Bytes" INTEGER, "UnitPrice" NUMERIC(10, 2) NOT NULL, PRIMARY KEY ("TrackId"), FOREIGN KEY("MediaTypeId") REFERENCES "MediaType" ("MediaTypeId"), FOREIGN KEY("GenreId") REFERENCES "Genre" ("GenreId"), FOREIGN KEY("AlbumId") REFERENCES "Album" ("AlbumId"))/*3 rows from Track table:TrackId Name AlbumId MediaTypeId GenreId Composer Milliseconds Bytes UnitPrice1 For Those About To Rock (We Salute You) 1 1 1 Angus Young, Malcolm Young, Brian Johnson 343719 11170334 0.992 Balls to the Wall 2 2 1 U. Dirkschneider, W. Hoffmann, H. Frank, P. Baltes, S. Kaufmann, G. Hoffmann 342562 5510424 0.993 Fast As a Shark 3 2 1 F. Baltes, S. Kaufman, U. Dirkscneider & W. Hoffman 230619 3990994 0.99*/================================== Ai Message ==================================Tool Calls: sql_db_query_checker (call_G64yYm6R6UauiVPCXJZMA49b) Call ID: call_G64yYm6R6UauiVPCXJZMA49b Args: query: SELECT Genre.Name, AVG(Track.Milliseconds) AS AverageLength FROM Track INNER JOIN Genre ON Track.GenreId = Genre.GenreId GROUP BY Genre.Name ORDER BY AverageLength DESC LIMIT 5;================================= Tool Message =================================Name: sql_db_query_checkerSELECT Genre.Name, AVG(Track.Milliseconds) AS AverageLength FROM Track INNER JOIN Genre ON Track.GenreId = Genre.GenreId GROUP BY Genre.Name ORDER BY AverageLength DESC LIMIT 5;================================== Ai Message ==================================Tool Calls: sql_db_query (call_AnO3SrhD0ODJBxh6dHMwvHwZ) Call ID: call_AnO3SrhD0ODJBxh6dHMwvHwZ Args: query: SELECT Genre.Name, AVG(Track.Milliseconds) AS AverageLength FROM Track INNER JOIN Genre ON Track.GenreId = Genre.GenreId GROUP BY Genre.Name ORDER BY AverageLength DESC LIMIT 5;================================= Tool Message =================================Name: sql_db_query[('Sci Fi & Fantasy', 2911783.0384615385), ('Science Fiction', 2625549.076923077), ('Drama', 2575283.78125), ('TV Shows', 2145041.0215053763), ('Comedy', 1585263.705882353)]================================== Ai Message ==================================On average, the genre with the longest tracks is "Sci Fi & Fantasy" with an average track length of approximately 2,911,783 milliseconds. This is followed by "Science Fiction," "Drama," "TV Shows," and "Comedy."
#sql_agent.py for studioimport pathlibfrom langchain.agents import create_agentfrom langchain.chat_models import init_chat_modelfrom langchain_community.agent_toolkits import SQLDatabaseToolkitfrom langchain_community.utilities import SQLDatabaseimport requests# Initialize an LLMmodel = init_chat_model("gpt-5.4")# Get the database, store it locallyurl = "https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db"local_path = pathlib.Path("Chinook.db")if local_path.exists(): print(f"{local_path} already exists, skipping download.")else: response = requests.get(url) if response.status_code == 200: local_path.write_bytes(response.content) print(f"File downloaded and saved as {local_path}") else: print(f"Failed to download the file. Status code: {response.status_code}")db = SQLDatabase.from_uri("sqlite:///Chinook.db")# Create the toolstoolkit = SQLDatabaseToolkit(db=db, llm=model)tools = toolkit.get_tools()for tool in tools: print(f"{tool.name}: {tool.description}\n")# Use create_agentsystem_prompt = """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 userspecifies a specific number of examples they wish to obtain, always limit yourquery to at most {top_k} results.You can order the results by a relevant column to return the most interestingexamples in the database. Never query for all the columns from a specific table,only ask for the relevant columns given the question.You MUST double check your query before executing it. If you get an error whileexecuting a query, rewrite the query and try again.DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to thedatabase.To start you should ALWAYS look at the tables in the database to see what youcan query. Do NOT skip this step.Then you should query the schema of the most relevant tables.""".format( dialect=db.dialect, top_k=5,)agent = create_agent( model, tools, system_prompt=system_prompt,)
question = "Which genre on average has the longest tracks?"config = {"configurable": {"thread_id": "1"}}for step in agent.stream( {"messages": [{"role": "user", "content": question}]}, config, stream_mode="values",): if "__interrupt__" in step: print("INTERRUPTED:") interrupt = step["__interrupt__"][0] for request in interrupt.value["action_requests"]: print(request["description"]) elif "messages" in step: step["messages"][-1].pretty_print() else: pass
...INTERRUPTED:Tool execution pending approvalTool: sql_db_queryArgs: {'query': 'SELECT g.Name AS Genre, AVG(t.Milliseconds) AS AvgTrackLength FROM Track t JOIN Genre g ON t.GenreId = g.GenreId GROUP BY g.Name ORDER BY AvgTrackLength DESC LIMIT 1;'}
from langgraph.types import Command for step in agent.stream( Command(resume={"decisions": [{"type": "approve"}]}), config, stream_mode="values",): if "messages" in step: step["messages"][-1].pretty_print() if "__interrupt__" in step: print("INTERRUPTED:") interrupt = step["__interrupt__"][0] for request in interrupt.value["action_requests"]: print(request["description"]) else: pass
================================== Ai Message ==================================Tool Calls: sql_db_query (call_7oz86Epg7lYRqi9rQHbZPS1U) Call ID: call_7oz86Epg7lYRqi9rQHbZPS1U Args: query: SELECT Genre.Name, AVG(Track.Milliseconds) AS AvgDuration FROM Track JOIN Genre ON Track.GenreId = Genre.GenreId GROUP BY Genre.Name ORDER BY AvgDuration DESC LIMIT 5;================================= Tool Message =================================Name: sql_db_query[('Sci Fi & Fantasy', 2911783.0384615385), ('Science Fiction', 2625549.076923077), ('Drama', 2575283.78125), ('TV Shows', 2145041.0215053763), ('Comedy', 1585263.705882353)]================================== Ai Message ==================================The genre with the longest average track length is "Sci Fi & Fantasy" with an average duration of about 2,911,783 milliseconds, followed by "Science Fiction" and "Drama."