from langchain_core.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import Runnable
from langchain_aws import ChatBedrock
import boto3
from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph.message import AnyMessage, add_messages
from langchain_core.messages import ToolMessage
from langchain_core.runnables import RunnableLambda
from langgraph.prebuilt import ToolNode
from langgraph.prebuilt import tools_condition

@tool
def compute_savings(monthly_cost: float) -> float:
    """
    根據(jù)用戶每月電費(fèi)計(jì)算切換到太陽能時(shí)潛在節(jié)省的工具。
    
    參數(shù)：
        monthly_cost (float)：用戶當(dāng)前的每月電費(fèi)。
    
    返回：
        dict：包含以下內(nèi)容的字典：
            - 'number_of_panels'：估計(jì)所需的太陽能板數(shù)量。
            - 'installation_cost'：估計(jì)的安裝成本。
            - 'net_savings_10_years'：安裝成本后的10年凈節(jié)省。
    """
    def calculate_solar_savings(monthly_cost):
        # 計(jì)算的假設(shè)
        cost_per_kWh = 0.28  
        cost_per_watt = 1.50  
        sunlight_hours_per_day = 3.5  
        panel_wattage = 350  
        system_lifetime_years = 10  

        # 每月用電量（千瓦時(shí)）
        monthly_consumption_kWh = monthly_cost / cost_per_kWh
        
        # 所需系統(tǒng)大小（千瓦）
        daily_energy_production = monthly_consumption_kWh / 30
        system_size_kW = daily_energy_production / sunlight_hours_per_day
        
        # 太陽能板數(shù)量和安裝成本
        number_of_panels = system_size_kW * 1000 / panel_wattage
        installation_cost = system_size_kW * 1000 * cost_per_watt
        
        # 年度和凈節(jié)省
        annual_savings = monthly_cost * 12
        total_savings_10_years = annual_savings * system_lifetime_years
        net_savings = total_savings_10_years - installation_cost
        
        return {
            "number_of_panels": round(number_of_panels),
            "installation_cost": round(installation_cost, 2),
            "net_savings_10_years": round(net_savings, 2)
        }

    # 返回計(jì)算的太陽能節(jié)省
    return calculate_solar_savings(monthly_cost)

def handle_tool_error(state) -> dict:
    """
    處理工具執(zhí)行期間發(fā)生的錯(cuò)誤的函數(shù)。
    
    參數(shù)：        state (dict)：AI智能體的當(dāng)前狀態(tài)，包括消息和工具調(diào)用詳情。
    
    返回：
        dict：包含每個(gè)遇到問題的工具的錯(cuò)誤消息的字典。
    """
    # 從當(dāng)前狀態(tài)中檢索錯(cuò)誤
    error = state.get("error")
    
    # 從狀態(tài)的消息歷史中獲取最后一個(gè)消息的工具調(diào)用
    tool_calls = state["messages"][-1].tool_calls
    
    # 返回包含錯(cuò)誤詳情的ToolMessages列表，與每個(gè)工具調(diào)用ID關(guān)聯(lián)
    return {
        "messages": [
            ToolMessage(
                cnotallow=f"錯(cuò)誤：{repr(error)}\n請(qǐng)修正你的錯(cuò)誤。",  # 為用戶格式化錯(cuò)誤消息
                tool_call_id=tc["id"],  # 將錯(cuò)誤消息與相應(yīng)的工具調(diào)用ID關(guān)聯(lián)
            )
            for tc in tool_calls  # 遍歷每個(gè)工具調(diào)用以產(chǎn)生單獨(dú)的錯(cuò)誤消息
        ]
    }

def create_tool_node_with_fallback(tools: list) -> dict:
    """
    創(chuàng)建具有后備錯(cuò)誤處理的工具節(jié)點(diǎn)的函數(shù)。
    
    參數(shù)：
        tools (list)：要包含在節(jié)點(diǎn)中的工具列表。
    
    返回：
        dict：如果發(fā)生錯(cuò)誤，使用后備行為的工具節(jié)點(diǎn)。
    """
    # 使用提供的工具創(chuàng)建ToolNode，并附加后備機(jī)制
    # 如果發(fā)生錯(cuò)誤，將調(diào)用handle_tool_error函數(shù)來管理錯(cuò)誤
    return ToolNode(tools).with_fallbacks(
        [RunnableLambda(handle_tool_error)],  # 使用lambda函數(shù)包裝錯(cuò)誤處理器
        exception_key="error"  # 指定這個(gè)后備是用于處理錯(cuò)誤的
    )

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

class Assistant:
    def __init__(self, runnable: Runnable):
        # 使用定義與工具交互過程的可運(yùn)行對(duì)象進(jìn)行初始化
        self.runnable = runnable

    def
 __call__(self, state: State):
        while True:
            # 使用當(dāng)前狀態(tài)（消息和上下文）調(diào)用可運(yùn)行對(duì)象
            result = self.runnable.invoke(state)
            
            # 如果工具未能返回有效輸出，重新提示用戶澄清或重試
            if not result.tool_calls and (
                not result.content
                or isinstance(result.content, list)
                and not result.content[0].get("text")
            ):
                # 添加請(qǐng)求有效響應(yīng)的消息
                messages = state["messages"] + [("user", "請(qǐng)給出一個(gè)真實(shí)的輸出。")]
                state = {**state, "messages": messages}
            else:
                # 當(dāng)獲得有效輸出時(shí)跳出循環(huán)
                break

        # 在處理完可運(yùn)行對(duì)象后返回最終狀態(tài)
        return {"messages": result}

def get_bedrock_client(region):
    return boto3.client("bedrock-runtime", region_name=region)

def create_bedrock_llm(client):
    return ChatBedrock(model_id='anthropic.claude-3-sonnet-20240229-v1:0', client=client, model_kwargs={'temperature': 0}, region_name='us-east-1')

llm = create_bedrock_llm(get_bedrock_client(reginotallow='us-east-1'))

primary_assistant_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            '''你是太陽能板比利時(shí)的樂于助人的客服助手。
            你應(yīng)該從他們那里獲取以下信息：
            - 每月電費(fèi)
            如果你不能清晰地識(shí)別這些信息，要求他們澄清！不要試圖胡亂猜測(cè)。

            在你能夠清晰地識(shí)別所有信息后，調(diào)用相關(guān)工具。
            ''',
        ),
        ("placeholder", "{messages}"),
    ]
)

# 定義助手將使用的工具
part_1_tools = [
    compute_savings
]

# 將工具綁定到助手的工作流程
part_1_assistant_runnable = primary_assistant_prompt | llm.bind_tools(part_1_tools)

AI智能體旨在計(jì)算太陽能板潛在節(jié)能節(jié)省。

builder = StateGraph(State)
builder.add_node("assistant", Assistant(part_1_assistant_runnable))
builder.add_node("tools", create_tool_node_with_fallback(part_1_tools))

邊定義了流程如何在節(jié)點(diǎn)之間移動(dòng)。在這里，助手開始對(duì)話，一旦收集到所需輸入，就過渡到工具，并在工具執(zhí)行后返回助手。

builder.add_edge(START, "assistant")  # 從助手開始
builder.add_conditional_edges("assistant", tools_condition)  # 輸入后移動(dòng)到工具
builder.add_edge("tools", "assistant")  # 工具執(zhí)行后返回助手

我們使用MemorySaver確保圖在不同步驟之間保留對(duì)話狀態(tài)。這允許助手記住用戶的輸入，確保多步驟交互中的連續(xù)性。

memory = MemorySaver()
graph = builder.compile(checkpointer=memory)

# import shutil
import uuid

# 讓我們創(chuàng)建一個(gè)用戶可能與助手進(jìn)行的示例對(duì)話
tutorial_questions = [
    '嘿',
    '你能計(jì)算我的節(jié)能嗎',
    “我的每月成本是100美元，我能節(jié)省多少”
]

thread_id = str(uuid.uuid4())

config = {
    "configurable": {
        "thread_id": thread_id,
    }
}

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