Spaces:
Sleeping
Sleeping
| from langchain_core.runnables import RunnableConfig | |
| from langgraph.prebuilt import tools_condition | |
| from langgraph.graph import END, StateGraph | |
| from langchain_core.messages import HumanMessage | |
| from src.langgraph.tools.scheduling_tools import ( | |
| create_a_activity, | |
| search_activities, | |
| update_a_activiy, | |
| delete_a_activity, | |
| ) | |
| from src.langgraph.langchain.llm import llm | |
| from src.langgraph.config.agent import Agent | |
| from src.langgraph.utils_function.function_graph import ( | |
| create_tool_node_with_fallback, | |
| human_review_node, | |
| ) | |
| from src.langgraph.langchain.prompt import ( | |
| scheduling_prompt, | |
| CompleteOrRoute, | |
| create_entry_node, | |
| ) | |
| from src.langgraph.state import State | |
| from src.utils.logger import logger | |
| from src.langgraph.tools.plan_itinerary import plan_itinerary | |
| class SchedulingAgent: | |
| def __init__(self, builder: StateGraph): | |
| self.builder = builder | |
| self.scheduling_safe_tools = [search_activities, plan_itinerary] | |
| self.scheduling_sensitive_tools = [ | |
| create_a_activity, | |
| update_a_activiy, | |
| delete_a_activity, | |
| ] | |
| self.scheduling_tools = ( | |
| self.scheduling_sensitive_tools + self.scheduling_safe_tools | |
| ) | |
| self.scheduling_runnable = scheduling_prompt | llm.bind_tools( | |
| self.scheduling_tools + [CompleteOrRoute] | |
| ) | |
| def route_scheduling(self, state: State): | |
| logger.info("Route scheduling") | |
| if ( | |
| state["messages_history"] is not None | |
| and state["messages"][0].content == "y" | |
| and "Do you want to run the following tool(s)?" | |
| in state["messages_history"][-1].content | |
| ): | |
| if state["messages"][-1].tool_calls and state["messages"][-1].tool_calls[0][ | |
| "name" | |
| ] in [r.name for r in self.scheduling_sensitive_tools]: | |
| logger.info("Sensitive tools") | |
| return "scheduling_sensitive_tools" | |
| logger.info("Safe tools") | |
| return END | |
| route = tools_condition(state) | |
| if route == END: | |
| return END | |
| tool_calls = state["messages"][-1].tool_calls | |
| did_cancel = any(tc["name"] == CompleteOrRoute.__name__ for tc in tool_calls) | |
| logger.info(f"Did cancel: {did_cancel}") | |
| if did_cancel: | |
| return "leave_skill" | |
| tool_names = [t.name for t in self.scheduling_safe_tools] | |
| if all(tc["name"] in tool_names for tc in tool_calls): | |
| logger.info("Scheduling safe tools") | |
| return "scheduling_safe_tools" | |
| logger.info("User review") | |
| return "user_review_scheduling" | |
| def node(self): | |
| self.builder.add_node( | |
| "enter_schedule_activity", | |
| create_entry_node("Schedule Activity Assistant", "scheduling"), | |
| ) | |
| self.builder.add_node("scheduling_agent", Agent(self.scheduling_runnable)) | |
| self.builder.add_edge("enter_schedule_activity", "scheduling_agent") | |
| self.builder.add_node( | |
| "scheduling_safe_tools", | |
| create_tool_node_with_fallback(self.scheduling_safe_tools), | |
| ) | |
| self.builder.add_node( | |
| "scheduling_sensitive_tools", | |
| create_tool_node_with_fallback(self.scheduling_sensitive_tools), | |
| ) | |
| self.builder.add_node("user_review_scheduling", human_review_node) | |
| def edge(self): | |
| self.builder.add_edge("enter_schedule_activity", "scheduling_agent") | |
| self.builder.add_conditional_edges( | |
| "scheduling_agent", | |
| self.route_scheduling, | |
| { | |
| "scheduling_sensitive_tools": "scheduling_sensitive_tools", | |
| END: "save_history", | |
| "leave_skill": "leave_skill", | |
| "scheduling_safe_tools": "scheduling_safe_tools", | |
| "user_review_scheduling": "user_review_scheduling", | |
| }, | |
| ) | |
| self.builder.add_edge("user_review_scheduling", "save_history") | |
| self.builder.add_edge("scheduling_sensitive_tools", "scheduling_agent") | |
| self.builder.add_edge("scheduling_safe_tools", "scheduling_agent") | |
| def __call__(self): | |
| self.node() | |
| self.edge() | |
| return self.builder | |