%%capture --no-stderr
% pip install -U langgraph langchain-community langchain-anthropic tavily-python pandas

import os
import requests
import sqlite3
import pandas as pd
import shutil

# 下載數(shù)據(jù)庫(kù)文件
db_url = "https://storage.googleapis.com/benchmarks-artifacts/travel-db/travel2.sqlite"
local_file = "travel2.sqlite"
backup_file = "travel2.backup.sqlite"
overwrite = False
if not os.path.exists(local_file) or overwrite:
    response = requests.get(db_url)
    response.raise_for_status()  # 確保請(qǐng)求成功
    with open(local_file, "wb") as file:
        file.write(response.content)

# 創(chuàng)建數(shù)據(jù)庫(kù)備份，以便在每個(gè)教程部分開始時(shí)重置數(shù)據(jù)庫(kù)狀態(tài)
shutil.copy(local_file, backup_file)

# 將航班數(shù)據(jù)更新為當(dāng)前時(shí)間，以適應(yīng)我們的教程
conn = sqlite3.connect(local_file)
cursor = conn.cursor()

# 讀取數(shù)據(jù)庫(kù)中的所有表
tables = pd.read_sql(
    "SELECT name FROM sqlite_master WHERE type='table';", conn
).name.tolist()
tdf = {}
for table_name in tables:
    tdf[table_name] = pd.read_sql(f"SELECT * from {table_name}", conn)

# 找到最早的出發(fā)時(shí)間，并計(jì)算時(shí)間差
example_time = pd.to_datetime(
    tdf["flights"]["actual_departure"].replace("\\N", pd.NaT)
).max()
current_time = pd.to_datetime("now").tz_localize(example_time.tz)
time_diff = current_time - example_time

# 更新預(yù)訂日期和航班時(shí)間
for column in ["book_date", "scheduled_departure", "scheduled_arrival", "actual_departure", "actual_arrival"]:
    tdf["flights"][column] = pd.to_datetime(
        tdf["flights"][column].replace("\\N", pd.NaT)
    ) + time_diff

# 將更新后的數(shù)據(jù)寫回?cái)?shù)據(jù)庫(kù)
for table_name, df in tdf.items():
    df.to_sql(table_name, conn, if_exists="replace", index=False)
conn.commit()
conn.close()

# 在本教程中，我們將使用這個(gè)本地文件作為數(shù)據(jù)庫(kù)
db = local_file

import re
import numpy as np
import openai
from langchain_core.tools import tool


@tool
def lookup_policy(query):
    """查詢公司政策，以確定某些選項(xiàng)是否允許。"""

from langchain_core.runnables import ensure_config
from typing import Optional
import sqlite3
import pytz
from datetime import datetime, timedelta, date

@tool
def fetch_user_flight_information():
    """獲取用戶的所有機(jī)票信息，包括航班詳情和座位分配。"""

@tool
def search_flights(
    departure_airport=None,
    arrival_airport=None,
    start_time=None,
    end_time=None,
    limit=20,
):
    """根據(jù)出發(fā)機(jī)場(chǎng)、到達(dá)機(jī)場(chǎng)和出發(fā)時(shí)間范圍來搜索航班。"""

@tool
def update_ticket_to_new_flight(ticket_no, new_flight_id):
    """將用戶的機(jī)票更新到一個(gè)新的有效航班上。"""

@tool
def cancel_ticket(ticket_no):
    """取消用戶的機(jī)票，并從數(shù)據(jù)庫(kù)中移除。"""

from typing import Optional, Union
from datetime import datetime, date

@tool
def search_car_rentals(
    locatinotallow=None,
    name=None,
    price_tier=None,
    start_date=None,
    end_date=None,
):
    """
    根據(jù)位置、公司名稱、價(jià)格等級(jí)、開始日期和結(jié)束日期來搜索租車服務(wù)。

    參數(shù)：
        location (Optional[str]): 租車服務(wù)的位置。
        name (Optional[str]): 租車公司的名稱。
        price_tier (Optional[str]): 租車的價(jià)格等級(jí)。
        start_date (Optional[Union[datetime, date]]): 租車的開始日期。
        end_date (Optional[Union[datetime, date]]): 租車的結(jié)束日期。

    返回：
        list[dict]: 匹配搜索條件的租車服務(wù)列表。
    """

@tool
def book_car_rental(rental_id):
    """
    通過租車ID來預(yù)訂租車服務(wù)。

    參數(shù)：
        rental_id (int): 要預(yù)訂的租車服務(wù)的ID。

    返回：
        str: 預(yù)訂成功與否的消息。
    """

@tool
def update_car_rental(
    rental_id,
    start_date=None,
    end_date=None,
):
    """
    通過租車ID來更新租車服務(wù)的開始和結(jié)束日期。

    參數(shù)：
        rental_id (int): 要更新的租車服務(wù)的ID。
        start_date (Optional[Union[datetime, date]]): 新的租車開始日期。
        end_date (Optional[Union[datetime, date]]): 新的租車結(jié)束日期。

    返回：
        str: 更新成功與否的消息。
    """

@tool
def cancel_car_rental(rental_id):
    """
    通過租車ID來取消租車服務(wù)。

    參數(shù)：
        rental_id (int): 要取消的租車服務(wù)的ID。

    返回：
        str: 取消成功與否的消息。
    """

@tool
def search_hotels(
    locatinotallow=None,
    name=None,
    price_tier=None,
    checkin_date=None,
    checkout_date=None,
):
    """
    根據(jù)位置、名稱、價(jià)格等級(jí)、入住日期和退房日期來搜索酒店。

    參數(shù)：
        location (Optional[str]): 酒店的位置。
        name (Optional[str]): 酒店的名稱。
        price_tier (Optional[str]): 酒店的價(jià)格等級(jí)。
        checkin_date
        
        # 入住日期和退房日期，用于搜索酒店
        checkin_date (Optional[Union[datetime, date]]): 酒店的入住日期。
        checkout_date (Optional[Union[datetime, date]]): 酒店的退房日期。

    返回：
        list[dict]: 符合搜索條件的酒店列表。
    """
    
@tool
def book_hotel(hotel_id):
    """
    通過酒店ID進(jìn)行預(yù)訂。

    參數(shù)：
        hotel_id (int): 要預(yù)訂的酒店的ID。

    返回：
        str: 預(yù)訂成功與否的消息。
    """
    
@tool
def update_hotel(
    hotel_id,
    checkin_date=None,
    checkout_date=None,
):
    """
    通過酒店ID更新酒店預(yù)訂的入住和退房日期。

    參數(shù)：
        hotel_id (int): 要更新預(yù)訂的酒店的ID。
        checkin_date (Optional[Union[datetime, date]]): 新的入住日期。
        checkout_date (Optional[Union[datetime, date]]): 新的退房日期。

    返回：
        str: 更新成功與否的消息。
    """
    
@tool
def cancel_hotel(hotel_id):
    """
    通過酒店ID取消酒店預(yù)訂。

    參數(shù)：
        hotel_id (int): 要取消預(yù)訂的酒店的ID。

    返回：
        str: 取消成功與否的消息。
    """

@tool
def search_trip_recommendations(
    locatinotallow=None,
    name=None,
    keywords=None,
):
    """
    根據(jù)位置、名稱和關(guān)鍵詞搜索旅行推薦。

    參數(shù)：
        location (Optional[str]): 旅行推薦的地點(diǎn)。
        name (Optional[str]): 旅行推薦的名字。
        keywords (Optional[str]): 與旅行推薦相關(guān)的關(guān)鍵詞。

    返回：
        list[dict]: 符合搜索條件的旅行推薦列表。
    """
    
@tool
def book_excursion(recommendation_id):
    """
    通過推薦ID預(yù)訂遠(yuǎn)足活動(dòng)。

    參數(shù)：
        recommendation_id (int): 要預(yù)訂的旅行推薦的ID。

    返回：
        str: 預(yù)訂成功與否的消息。
    """
    
@tool
def update_excursion(recommendation_id, details):
    """
    通過推薦ID更新旅行推薦的細(xì)節(jié)。

    參數(shù)：
        recommendation_id (int): 要更新的旅行推薦的ID。
        details (str): 旅行推薦的新細(xì)節(jié)。

    返回：
        str: 更新成功與否的消息。
    """
    
@tool
def cancel_excursion(recommendation_id):
    """
    通過推薦ID取消旅行推薦。

    參數(shù)：
        recommendation_id (int): 要取消的旅行推薦的ID。

    返回：
        str: 取消成功與否的消息。
    """

from langgraph.prebuilt import ToolNode
from langchain_core.runnables import RunnableLambda

def handle_tool_error(state):
    error = state.get("error")
    tool_calls = state["messages"][-1].tool_calls
    return {
        "messages": [
            ToolMessage(
                cnotallow=f"錯(cuò)誤: {repr(error)}\n請(qǐng)修正你的錯(cuò)誤。",
                tool_call_id=tc["id"],
            )
            for tc in tool_calls
        ]
    }

def create_tool_node_with_fallback(tools):
    return ToolNode(tools).with_fallbacks(
        [RunnableLambda(handle_tool_error)], exception_key="error"
    )

def _print_event(event, _printed, max_length=1500):
    current_state = event.get("dialog_state")
    if current_state:
        print(f"當(dāng)前狀態(tài): ", current_state[-1])
    message = event.get("messages")
    if message:
        if isinstance(message, list):
            message = message[-1]
        if message.id not in _printed:
            msg_repr = message.pretty_repr(html=True)
            if len(msg_repr) > max_length:
                msg_repr = msg_repr[:max_length] + " ... （內(nèi)容已截?cái)啵?
            print(msg_repr)
            _printed.add(message.id)

from langgraph.graph.message import add_messages, AnyMessage
from typing_extensions import TypedDict
from typing import Annotated


class State(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]

from langchain_core.runnables import Runnable, RunnableConfig
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_anthropic import ChatAnthropic
from langchain_core.prompts import ChatPromptTemplate


class Assistant:
    def __init__(self, runnable: Runnable):
        self.runnable = runnable

    def __call__(self, state: State, config: RunnableConfig):
        while True:
            passenger_id = config.get("passenger_id", None)
            state = {**state, "user_info": passenger_id}
            result = self.runnable.invoke(state)
            # 如果大型語(yǔ)言模型返回了一個(gè)空響應(yīng)，我們將重新提示它給出一個(gè)實(shí)際的響應(yīng)。
            if (
                not result.content
                or isinstance(result.content, list)
                and not result.content[0].get("text")
            ):
                messages = state["messages"] + [("user", "請(qǐng)給出一個(gè)真實(shí)的輸出。")]
                state = {**state, "messages": messages}
            else:
                break
        return {"messages": result}


# Haiku模型更快、成本更低，但準(zhǔn)確性稍差
# llm = ChatAnthropic(model="claude-3-haiku-20240307")
llm = ChatAnthropic(model="claude-3-sonnet-20240229", temperature=1)

primary_assistant_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "你是一個(gè)為瑞士航空提供幫助的客戶支持助手。"
            "使用提供的工具來搜索航班、公司政策和其他信息以幫助回答用戶的查詢。"
            "在搜索時(shí)，要有毅力。如果第一次搜索沒有結(jié)果，就擴(kuò)大你的查詢范圍。"
            "如果搜索結(jié)果為空，不要放棄，先擴(kuò)大搜索范圍。"
            "\n\n當(dāng)前用戶：\n<User>\n{user_info}\n</User>"
            "\n當(dāng)前時(shí)間：{time}。",
        ),
        ("placeholder", "{messages}"),
    ]
).partial(time=datetime.now())

part_1_tools = [
    TavilySearchResults(max_results=1),
    fetch_user_flight_information,
    search_flights,
    lookup_policy,
    update_ticket_to_new_flight,
    cancel_ticket,
    search_car_rentals,
    book_car_rental,
    update_car_rental,
    cancel_car_rental,
    search_hotels,
    book_hotel,
    update_hotel,
    cancel_hotel,
    search_trip_recommendations,
    book_excursion,
    update_excursion,
    cancel_excursion,
]
part_1_assistant_runnable = primary_assistant_prompt | llm.bind_tools(part_1_tools)

from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import tools_condition, ToolNode

builder = StateGraph(State)


# 定義節(jié)點(diǎn)：這些節(jié)點(diǎn)執(zhí)行具體的工作
builder.add_node("assistant", Assistant(part_1_assistant_runnable))
builder.add_node("action", create_tool_node_with_fallback(part_1_tools))
# 定義邊：這些邊決定了控制流程如何移動(dòng)
builder.set_entry_point("assistant")
builder.add_conditional_edges(
    "assistant",
    tools_condition,
    # "action"調(diào)用我們的工具之一。END導(dǎo)致圖終止（并向用戶做出響應(yīng)）
    {"action": "action", END: END},
)
builder.add_edge("action", "assistant")

# 檢查點(diǎn)器允許圖保存其狀態(tài)
# 這是整個(gè)圖的完整記憶。
memory = SqliteSaver.from_conn_string(":memory:")
part_1_graph = builder.compile(checkpointer=memory)

from IPython.display import Image, display

try:
    display(Image(part_1_graph.get_graph(xray=True).draw_mermaid_png()))
except:
    # 這需要一些額外的依賴項(xiàng)，是可選的
    pass

import uuid
import shutil

# 假設(shè)這是用戶與助手之間可能發(fā)生的對(duì)話示例
tutorial_questions = [
    "你好，我的航班是什么時(shí)候？",
    "我可以把我的航班改簽到更早的時(shí)間嗎？我想今天晚些時(shí)候離開。",
    "那就把我的航班改簽到下周某個(gè)時(shí)間吧",
    "下一個(gè)可用的選項(xiàng)很好",
    "住宿和交通方面有什么建議？",
    "我想在為期一周的住宿中選擇一個(gè)經(jīng)濟(jì)實(shí)惠的酒店（7天），并且我還想租一輛車。",
    "好的，你能為你推薦的酒店預(yù)訂嗎？聽起來不錯(cuò)。",
    "是的，去預(yù)訂任何中等價(jià)位且有可用性的酒店。",
    "對(duì)于汽車，我有哪些選擇？",
    "太棒了，我們只選擇最便宜的選項(xiàng)。預(yù)訂7天。",
    "那么，你對(duì)我的旅行有什么建議？",
    "在我在那里的時(shí)候，有哪些活動(dòng)是可用的？",
    "有趣 - 我喜歡博物館，有哪些選擇？",
    "好的，那就為我在那里的第二天預(yù)訂一個(gè)。",
]

# 使用備份文件以便我們可以從每個(gè)部分的原始位置重新啟動(dòng)
shutil.copy(backup_file, db)
thread_id = str(uuid.uuid4())

config = {
    "configurable": {
        # passenger_id 在我們的航班工具中使用
        # 以獲取用戶的航班信息
        "passenger_id": "3442 587242",
        # 檢查點(diǎn)通過 thread_id 訪問
        "thread_id": thread_id,
    }
}


_printed = set()
for question in tutorial_questions:
    events = part_1_graph.stream(
        {"messages": ("user", question)}, config, stream_mode="values"
    )
    for event in events:
        _print_event(event, _printed)