from langchain_core.tools import toolfrom langchain_community.tools.tavily_search import TavilySearchResultsfrom langchain_core.prompts import ChatPromptTemplatefrom langchain_core.runnables import Runnablefrom langchain_aws import ChatBedrockimport boto3from typing import Annotatedfrom typing_extensions import TypedDictfrom langgraph.graph.message import AnyMessage, add_messagesfrom langchain_core.messages import ToolMessagefrom langchain_core.runnables import RunnableLambdafrom langgraph.prebuilt import ToolNode

@tooldef compute_savings(monthly_cost: float) -> float:"""Tool to compute the potential savings when switching to solar energy based on the user's monthly electricity cost.Args:monthly_cost (float): The user's current monthly electricity cost.Returns:dict: A dictionary containing:- 'number_of_panels': The estimated number of solar panels required.- 'installation_cost': The estimated installation cost.- 'net_savings_10_years': The net savings over 10 years after installation costs."""def calculate_solar_savings(monthly_cost):# Assumptions for the calculationcost_per_kWh = 0.28cost_per_watt = 1.50sunlight_hours_per_day = 3.5panel_wattage = 350system_lifetime_years = 10
# Monthly electricity consumption in kWhmonthly_consumption_kWh = monthly_cost / cost_per_kWh# Required system size in kWdaily_energy_production = monthly_consumption_kWh / 30system_size_kW = daily_energy_production / sunlight_hours_per_day# Number of panels and installation costnumber_of_panels = system_size_kW * 1000 / panel_wattageinstallation_cost = system_size_kW * 1000 * cost_per_watt# Annual and net savingsannual_savings = monthly_cost * 12total_savings_10_years = annual_savings * system_lifetime_yearsnet_savings = total_savings_10_years - installation_costreturn {"number_of_panels": round(number_of_panels),"installation_cost": round(installation_cost, 2),"net_savings_10_years": round(net_savings, 2)}
# Return calculated solar savingsreturn calculate_solar_savings(monthly_cost)

def handle_tool_error(state) -> dict:"""Function to handle errors that occur during tool execution.Args:state (dict): The current state of the AI agent, which includes messages and tool call details.Returns:dict: A dictionary containing error messages for each tool that encountered an issue."""# Retrieve the error from the current stateerror = state.get("error")# Access the tool calls from the last message in the state's message historytool_calls = state["messages"][-1].tool_calls# Return a list of ToolMessages with error details, linked to each tool call IDreturn {"messages": [ToolMessage(content=f"Error: {repr(error)}\n please fix your mistakes.",# Format the error message for the usertool_call_id=tc["id"],# Associate the error message with the corresponding tool call ID)for tc in tool_calls# Iterate over each tool call to produce individual error messages]}
def create_tool_node_with_fallback(tools: list) -> dict:"""Function to create a tool node with fallback error handling.Args:tools (list): A list of tools to be included in the node.Returns:dict: A tool node that uses fallback behavior in case of errors."""# Create a ToolNode with the provided tools and attach a fallback mechanism# If an error occurs, it will invoke the handle_tool_error function to manage the errorreturn ToolNode(tools).with_fallbacks([RunnableLambda(handle_tool_error)],# Use a lambda function to wrap the error handlerexception_key="error"# Specify that this fallback is for handling errors)

class State(TypedDict):messages: Annotated[list[AnyMessage], add_messages]
class Assistant:def __init__(self, runnable: Runnable):# Initialize with the runnable that defines the process for interacting with the toolsself.runnable = runnable
def __call__(self, state: State):while True:# Invoke the runnable with the current state (messages and context)result = self.runnable.invoke(state)# If the tool fails to return valid output, re-prompt the user to clarify or retryif not result.tool_calls and (not result.contentor isinstance(result.content, list)and not result.content[0].get("text")):# Add a message to request a valid responsemessages = state["messages"] + [("user", "Respond with a real output.")]state = {**state, "messages": messages}else:# Break the loop when valid output is obtainedbreak
# Return the final state after processing the runnablereturn {"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(region='us-east-1'))

primary_assistant_prompt = ChatPromptTemplate.from_messages([("system",'''You are a helpful customer support assistant for Solar Panels Belgium.You should get the following information from them:- monthly electricity costIf you are not able to discern this info, ask them to clarify! Do not attempt to wildly guess.
After you are able to discern all the information, call the relevant tool.''',),("placeholder", "{messages}"),])

builder = StateGraph(State)builder.add_node("assistant", Assistant(part_1_assistant_runnable))builder.add_node("tools", create_tool_node_with_fallback(part_1_tools))builder.add_edge(START, "assistant")# Start with the assistantbuilder.add_conditional_edges("assistant", tools_condition)# Move to tools after inputbuilder.add_edge("tools", "assistant")# Return to assistant after tool executio
memory = MemorySaver()graph = builder.compile(checkpointer=memory)

# import shutilimport uuid
# Let's create an example conversation a user might have with the assistanttutorial_questions = ['hey','can you calculate my energy saving',"my montly cost is $100, what will i save"]
# Update with the backup file so we can restart from the original place in each section# shutil.copy(backup_file, db)thread_id = str(uuid.uuid4())
config = {"configurable": {# The passenger_id is used in our flight tools to# fetch the user's flight information# "passenger_id": "3442 587242",# Checkpoints are accessed by thread_id"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)