mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 06:35:46 +02:00
59 KiB
59 KiB
In [32]:
import os
In [33]:
# import requests
#
# url = "https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db"
#
# response = requests.get(url)
#
# if response.status_code == 200:
# # Open a local file in binary write mode
# with open("Chinook.db", "wb") as file:
# # Write the content of the response (the file) to the local file
# file.write(response.content)
# print("File downloaded and saved as Chinook.db")
# else:
# print(f"Failed to download the file. Status code: {response.status_code}")File downloaded and saved as Chinook.db
In [34]:
%%capture --no-stderr --no-display
!pip install langchain_community langchain_openaiIn [84]:
from langchain_community.utilities import SQLDatabase
db = SQLDatabase.from_uri("sqlite:///Chinook.db")
print(db.dialect)
print(db.get_usable_table_names())
db.run("SELECT * FROM Artist LIMIT 10;")Out [84]:
sqlite ['Album', 'Artist', 'Customer', 'Employee', 'Genre', 'Invoice', 'InvoiceLine', 'MediaType', 'Playlist', 'PlaylistTrack', 'Track']
"[(1, 'AC/DC'), (2, 'Accept'), (3, 'Aerosmith'), (4, 'Alanis Morissette'), (5, 'Alice In Chains'), (6, 'Antônio Carlos Jobim'), (7, 'Apocalyptica'), (8, 'Audioslave'), (9, 'BackBeat'), (10, 'Billy Cobham')]"
In [36]:
from typing import Any, Dict, List
from langchain_core.runnables import RunnableLambda, RunnableWithFallbacks
from langchain_core.messages import ToolMessage, AIMessage
from langgraph.prebuilt import ToolNode
def create_tool_node_with_fallback(tools: list) -> RunnableWithFallbacks[Any, dict]:
"""
Create a ToolNode with a fallback to handle errors and surface them to the agent.
"""
return ToolNode(tools).with_fallbacks(
[RunnableLambda(handle_tool_error)], exception_key="error"
)
def handle_tool_error(state) -> dict:
error = state.get("error")
tool_calls = state["messages"][-1].tool_calls
return {
"messages": [
ToolMessage(
content=f"Error: {repr(error)}\n please fix your mistakes.",
tool_call_id=tc["id"],
)
for tc in tool_calls
]
}In [37]:
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_openai import ChatOpenAI
toolkit = SQLDatabaseToolkit(db=db, llm=ChatOpenAI(model="gpt-4o"))
tools = toolkit.get_tools()
list_tables_tool = next(tool for tool in tools if tool.name == "sql_db_list_tables")
get_schema_tool = next(tool for tool in tools if tool.name == "sql_db_schema")
print(list_tables_tool(""))
print(get_schema_tool("Artist"))Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track
CREATE TABLE "Artist" (
"ArtistId" INTEGER NOT NULL,
"Name" NVARCHAR(120),
PRIMARY KEY ("ArtistId")
)
/*
3 rows from Artist table:
ArtistId Name
1 AC/DC
2 Accept
3 Aerosmith
*/
In [80]:
from langchain.agents import tool
from langchain_core.prompts import ChatPromptTemplate
query_check_system = """You are a SQL expert with a strong attention to detail.
Double check the SQLite 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."""
query_check_prompt = ChatPromptTemplate.from_messages([("system", query_check_system),("user", "{query}")])
query_check = query_check_prompt | ChatOpenAI(model="gpt-4o", temperature=0)
@tool
def check_query_tool(query: str) -> str:
"""
Use this tool to double-check if your query is correct before executing it.
"""
return query_check.invoke({"query": query}).content
@tool
def db_query_tool(query: str) -> str:
"""
Execute a SQL query against the database and get back the result.
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.
"""
result = db.run_no_throw(query)
if not result:
return "Error: Query failed. Please rewrite your query and try again."
return result
print(check_query_tool("SELET * FROM Artist LIMIT 10;"))
print(db_query_tool("SELECT * FROM Artist LIMIT 10;"))There is a typo in the SQL keyword. The correct keyword is `SELECT` instead of `SELET`. Here is the corrected query: ```sql SELECT * FROM Artist LIMIT 10; ``` [(1, 'AC/DC'), (2, 'Accept'), (3, 'Aerosmith'), (4, 'Alanis Morissette'), (5, 'Alice In Chains'), (6, 'Antônio Carlos Jobim'), (7, 'Apocalyptica'), (8, 'Audioslave'), (9, 'BackBeat'), (10, 'Billy Cobham')]
In [81]:
from typing import Annotated
from typing_extensions import TypedDict
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph
from langgraph.graph.message import AnyMessage, add_messages
from langgraph.prebuilt.tool_node import ToolNode
from langchain_core.messages import AIMessage
from typing import Literal
# Define the state for the agent
class State(TypedDict):
messages: Annotated[list[AnyMessage], add_messages]
# Define a new graph
workflow = StateGraph(State)
# Add a node for the first tool call
def first_tool_call(state: State) -> dict[str, list[AIMessage]]:
return {
"messages": [
AIMessage(
content="",
tool_calls=[
{
"name": "sql_db_list_tables",
"args": {
},
"id": "tool_abcd123",
}
],
)
]
}
workflow.add_node("first_tool_call", first_tool_call)
# Add nodes for the first two tools
workflow.add_node("list_tables_tool", create_tool_node_with_fallback([list_tables_tool]))
workflow.add_node("get_schema_tool", create_tool_node_with_fallback([get_schema_tool]))
# Add a node for a model to choose the relevant tables based on the question and available tables
model_get_schema = ChatOpenAI(model="gpt-4o", temperature=0).bind_tools([get_schema_tool])
workflow.add_node("model_get_schema", lambda state: {"messages": [model_get_schema.invoke(state["messages"])],})
# Add a node for a model to generate a query based on the question and schema
query_gen_system = """You are a SQL expert with a strong attention to detail.
Given an input question, create a syntactically correct SQLite 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 5 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.
If you get an error while executing a query, rewrite the query and try again.
If you get an empty result set, you should try to rewrite the query to get a non-empty result set.
NEVER make stuff up if you don't have enough information to answer the query.
If you are unsure about your query, you should check it with the appropriate tool.
If you have enough information to answer the input question, simply reply with the final answer.
DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database."""
query_gen_prompt = ChatPromptTemplate.from_messages([("system", query_gen_system),("placeholder", "{messages}")])
query_gen = query_gen_prompt | ChatOpenAI(model="gpt-4o", temperature=0).bind_tools([db_query_tool, check_query_tool])
workflow.add_node("query_gen", lambda state: {"messages": [query_gen.invoke(state)],})
# Add nodes for the last two tools, check_query_tool and db_query_tool
workflow.add_node("sql_actions", create_tool_node_with_fallback([check_query_tool, db_query_tool]))
# Define a conditional edge to decide whether to continue or end the workflow
def should_continue(state: State) -> Literal[END, "sql_actions"]:
messages = state["messages"]
last_message = messages[-1]
# If there is no tool call, then we finish
if not last_message.tool_calls:
return END
else:
return "sql_actions"
# Specify the edges between the nodes
workflow.set_entry_point("first_tool_call")
workflow.add_edge("first_tool_call", "list_tables_tool")
workflow.add_edge("list_tables_tool", "model_get_schema")
workflow.add_edge("model_get_schema", "get_schema_tool")
workflow.add_edge("get_schema_tool", "query_gen")
workflow.add_conditional_edges(
"query_gen",
should_continue,
)
workflow.add_edge("sql_actions", "query_gen")
# Compile the workflow into a runnable
app = workflow.compile()In [82]:
from langchain_core.runnables.graph import MermaidDrawMethod
from IPython.display import display, Image
display(
Image(
app.get_graph().draw_mermaid_png(
draw_method=MermaidDrawMethod.API,
)
)
)
In [88]:
for event in app.stream({"messages": [("user", "Which sales agent made the most in sales in 2009?")]}):
print(event){'first_tool_call': {'messages': [AIMessage(content='', tool_calls=[{'name': 'sql_db_list_tables', 'args': {}, 'id': 'tool_abcd123'}])]}}
{'list_tables_tool': {'messages': [ToolMessage(content='Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track', name='sql_db_list_tables', tool_call_id='tool_abcd123')]}}
{'model_get_schema': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_18NT4ma6CkCYn1r4evZ2qgAx', 'function': {'arguments': '{"table_names":"Employee, Invoice"}', 'name': 'sql_db_schema'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 18, 'prompt_tokens': 177, 'total_tokens': 195}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_319be4768e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-7f3fc204-b27b-4a40-8921-faaa2d4e45b6-0', tool_calls=[{'name': 'sql_db_schema', 'args': {'table_names': 'Employee, Invoice'}, 'id': 'call_18NT4ma6CkCYn1r4evZ2qgAx'}], usage_metadata={'input_tokens': 177, 'output_tokens': 18, 'total_tokens': 195})]}}
{'get_schema_tool': {'messages': [ToolMessage(content='\nCREATE TABLE "Employee" (\n\t"EmployeeId" INTEGER NOT NULL, \n\t"LastName" NVARCHAR(20) NOT NULL, \n\t"FirstName" NVARCHAR(20) NOT NULL, \n\t"Title" NVARCHAR(30), \n\t"ReportsTo" INTEGER, \n\t"BirthDate" DATETIME, \n\t"HireDate" DATETIME, \n\t"Address" NVARCHAR(70), \n\t"City" NVARCHAR(40), \n\t"State" NVARCHAR(40), \n\t"Country" NVARCHAR(40), \n\t"PostalCode" NVARCHAR(10), \n\t"Phone" NVARCHAR(24), \n\t"Fax" NVARCHAR(24), \n\t"Email" NVARCHAR(60), \n\tPRIMARY KEY ("EmployeeId"), \n\tFOREIGN KEY("ReportsTo") REFERENCES "Employee" ("EmployeeId")\n)\n\n/*\n3 rows from Employee table:\nEmployeeId\tLastName\tFirstName\tTitle\tReportsTo\tBirthDate\tHireDate\tAddress\tCity\tState\tCountry\tPostalCode\tPhone\tFax\tEmail\n1\tAdams\tAndrew\tGeneral Manager\tNone\t1962-02-18 00:00:00\t2002-08-14 00:00:00\t11120 Jasper Ave NW\tEdmonton\tAB\tCanada\tT5K 2N1\t+1 (780) 428-9482\t+1 (780) 428-3457\tandrew@chinookcorp.com\n2\tEdwards\tNancy\tSales Manager\t1\t1958-12-08 00:00:00\t2002-05-01 00:00:00\t825 8 Ave SW\tCalgary\tAB\tCanada\tT2P 2T3\t+1 (403) 262-3443\t+1 (403) 262-3322\tnancy@chinookcorp.com\n3\tPeacock\tJane\tSales Support Agent\t2\t1973-08-29 00:00:00\t2002-04-01 00:00:00\t1111 6 Ave SW\tCalgary\tAB\tCanada\tT2P 5M5\t+1 (403) 262-3443\t+1 (403) 262-6712\tjane@chinookcorp.com\n*/\n\n\nCREATE TABLE "Invoice" (\n\t"InvoiceId" INTEGER NOT NULL, \n\t"CustomerId" INTEGER NOT NULL, \n\t"InvoiceDate" DATETIME NOT NULL, \n\t"BillingAddress" NVARCHAR(70), \n\t"BillingCity" NVARCHAR(40), \n\t"BillingState" NVARCHAR(40), \n\t"BillingCountry" NVARCHAR(40), \n\t"BillingPostalCode" NVARCHAR(10), \n\t"Total" NUMERIC(10, 2) NOT NULL, \n\tPRIMARY KEY ("InvoiceId"), \n\tFOREIGN KEY("CustomerId") REFERENCES "Customer" ("CustomerId")\n)\n\n/*\n3 rows from Invoice table:\nInvoiceId\tCustomerId\tInvoiceDate\tBillingAddress\tBillingCity\tBillingState\tBillingCountry\tBillingPostalCode\tTotal\n1\t2\t2021-01-01 00:00:00\tTheodor-Heuss-Straße 34\tStuttgart\tNone\tGermany\t70174\t1.98\n2\t4\t2021-01-02 00:00:00\tUllevålsveien 14\tOslo\tNone\tNorway\t0171\t3.96\n3\t8\t2021-01-03 00:00:00\tGrétrystraat 63\tBrussels\tNone\tBelgium\t1000\t5.94\n*/', name='sql_db_schema', tool_call_id='call_18NT4ma6CkCYn1r4evZ2qgAx')]}}
{'query_gen': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_9XavnPA7E11gx3QcWFe2qx32', 'function': {'arguments': '{"table_names":"Customer"}', 'name': 'sql_db_schema'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 16, 'prompt_tokens': 1200, 'total_tokens': 1216}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_319be4768e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-00ee8782-f824-4656-b1d6-0e93183f2db7-0', tool_calls=[{'name': 'sql_db_schema', 'args': {'table_names': 'Customer'}, 'id': 'call_9XavnPA7E11gx3QcWFe2qx32'}], usage_metadata={'input_tokens': 1200, 'output_tokens': 16, 'total_tokens': 1216})]}}
{'sql_actions': {'messages': [ToolMessage(content="Error: KeyError('sql_db_schema')\n please fix your mistakes.", tool_call_id='call_9XavnPA7E11gx3QcWFe2qx32')]}}
{'query_gen': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_C4ROOfoerwPv2b8lq9LqQnD2', 'function': {'arguments': '{"query":"SELECT e.FirstName, e.LastName, SUM(i.Total) as TotalSales FROM Employee e JOIN Customer c ON e.EmployeeId = c.SupportRepId JOIN Invoice i ON c.CustomerId = i.CustomerId WHERE strftime(\'%Y\', i.InvoiceDate) = \'2009\' GROUP BY e.EmployeeId ORDER BY TotalSales DESC LIMIT 5;"}', 'name': 'check_query_tool'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 85, 'prompt_tokens': 1239, 'total_tokens': 1324}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_319be4768e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-37006782-d3f3-4a77-873a-19bcd9d39cfd-0', tool_calls=[{'name': 'check_query_tool', 'args': {'query': "SELECT e.FirstName, e.LastName, SUM(i.Total) as TotalSales FROM Employee e JOIN Customer c ON e.EmployeeId = c.SupportRepId JOIN Invoice i ON c.CustomerId = i.CustomerId WHERE strftime('%Y', i.InvoiceDate) = '2009' GROUP BY e.EmployeeId ORDER BY TotalSales DESC LIMIT 5;"}, 'id': 'call_C4ROOfoerwPv2b8lq9LqQnD2'}], usage_metadata={'input_tokens': 1239, 'output_tokens': 85, 'total_tokens': 1324})]}}
{'sql_actions': {'messages': [ToolMessage(content="The provided query looks mostly correct, but let's double-check for common mistakes:\n\n1. **Using NOT IN with NULL values**: This query does not use `NOT IN`.\n2. **Using UNION when UNION ALL should have been used**: This query does not use `UNION`.\n3. **Using BETWEEN for exclusive ranges**: This query does not use `BETWEEN`.\n4. **Data type mismatch in predicates**: The `strftime` function returns a string, and the comparison to `'2009'` is correct.\n5. **Properly quoting identifiers**: Identifiers are not quoted, but they do not contain any special characters or reserved words, so this is acceptable.\n6. **Using the correct number of arguments for functions**: The `strftime` function is used correctly with the right number of arguments.\n7. **Casting to the correct data type**: No explicit casting is required here.\n8. **Using the proper columns for joins**: The joins are correctly using `EmployeeId`, `SupportRepId`, and `CustomerId`.\n\nSince there are no mistakes, the original query is correct. Here it is again:\n\n```sql\nSELECT e.FirstName, e.LastName, SUM(i.Total) as TotalSales \nFROM Employee e \nJOIN Customer c ON e.EmployeeId = c.SupportRepId \nJOIN Invoice i ON c.CustomerId = i.CustomerId \nWHERE strftime('%Y', i.InvoiceDate) = '2009' \nGROUP BY e.EmployeeId \nORDER BY TotalSales DESC \nLIMIT 5;\n```", name='check_query_tool', tool_call_id='call_C4ROOfoerwPv2b8lq9LqQnD2')]}}
{'query_gen': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_I2wGxopD5yAx5I0ghlx7dM4k', 'function': {'arguments': '{"query":"SELECT e.FirstName, e.LastName, SUM(i.Total) as TotalSales FROM Employee e JOIN Customer c ON e.EmployeeId = c.SupportRepId JOIN Invoice i ON c.CustomerId = i.CustomerId WHERE strftime(\'%Y\', i.InvoiceDate) = \'2009\' GROUP BY e.EmployeeId ORDER BY TotalSales DESC LIMIT 5;"}', 'name': 'db_query_tool'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 85, 'prompt_tokens': 1648, 'total_tokens': 1733}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_319be4768e', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-9751c70f-a1e6-454e-86db-f59bb856c0ce-0', tool_calls=[{'name': 'db_query_tool', 'args': {'query': "SELECT e.FirstName, e.LastName, SUM(i.Total) as TotalSales FROM Employee e JOIN Customer c ON e.EmployeeId = c.SupportRepId JOIN Invoice i ON c.CustomerId = i.CustomerId WHERE strftime('%Y', i.InvoiceDate) = '2009' GROUP BY e.EmployeeId ORDER BY TotalSales DESC LIMIT 5;"}, 'id': 'call_I2wGxopD5yAx5I0ghlx7dM4k'}], usage_metadata={'input_tokens': 1648, 'output_tokens': 85, 'total_tokens': 1733})]}}
{'sql_actions': {'messages': [ToolMessage(content="[('Steve', 'Johnson', 164.34), ('Margaret', 'Park', 161.37), ('Jane', 'Peacock', 123.75)]", name='db_query_tool', tool_call_id='call_I2wGxopD5yAx5I0ghlx7dM4k')]}}
{'query_gen': {'messages': [AIMessage(content='The sales agent who made the most in sales in 2009 is Steve Johnson, with total sales of 164.34.', response_metadata={'token_usage': {'completion_tokens': 27, 'prompt_tokens': 1778, 'total_tokens': 1805}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_319be4768e', 'finish_reason': 'stop', 'logprobs': None}, id='run-07c3b283-e542-45df-ac4e-3165de5eda11-0', usage_metadata={'input_tokens': 1778, 'output_tokens': 27, 'total_tokens': 1805})]}}
In [98]:
db_query_tool.invoke("SELECT * FROM Artist LIMIT 10;")Out [98]:
"[(1, 'AC/DC'), (2, 'Accept'), (3, 'Aerosmith'), (4, 'Alanis Morissette'), (5, 'Alice In Chains'), (6, 'Antônio Carlos Jobim'), (7, 'Apocalyptica'), (8, 'Audioslave'), (9, 'BackBeat'), (10, 'Billy Cobham')]"
In [108]:
from datetime import datetime, timezone
# The given timestamp
timestamp = "2024-06-11T09:50:30.784939"
ts2 = "2024-06-11T07:37:48.308394+00:00"
# Parse the timestamp into a datetime object
dt = datetime.fromisoformat(ts2)
# Print the resulting datetime object
print(dt)
dt.tzinfoOut [108]:
2024-06-11 07:37:48.308394+00:00
datetime.timezone.utc
In [109]:
print(dt.tzinfo)UTC
In [106]:
dt.replace(tzinfo=timezone.utc)Out [106]:
datetime.datetime(2024, 6, 11, 9, 50, 30, 784939, tzinfo=datetime.timezone.utc)
In [110]:
dtOut [110]:
datetime.datetime(2024, 6, 11, 7, 37, 48, 308394, tzinfo=datetime.timezone.utc)
In [ ]: