Spaces:
Sleeping
Sleeping
| from langgraph.graph import StateGraph, END, MessagesState | |
| from langgraph.prebuilt import tools_condition,ToolNode | |
| from langchain_core.prompts import ChatPromptTemplate | |
| import datetime | |
| from langgraph.types import interrupt, Command | |
| from langgraph.checkpoint.memory import MemorySaver | |
| import streamlit as st | |
| #module import | |
| from src.langgraphagenticai.node.sdlc_node import SDLCNode | |
| from src.langgraphagenticai.node.ai_news_node import AINewsNode | |
| from src.langgraphagenticai.node import travel_planner_node | |
| from src.langgraphagenticai.node.customer_support_chatbot import Customer_Support_Bot | |
| from src.langgraphagenticai.tools.customtool import book_appointment, cancel_appointment, get_next_available_appointment | |
| from src.langgraphagenticai.tools.search_tool import create_tool_node, get_tools | |
| from src.langgraphagenticai.node.chatbot_with_tool_node import ChatbotWithToolNode | |
| from src.langgraphagenticai.node.basic_chatbot_node import BasicChatbotNode | |
| from src.langgraphagenticai.state.state import State , SDLCState | |
| from src.langgraphagenticai.node.travel_planner_node import TravelPlannerNode | |
| class GraphBuilder: | |
| """ | |
| Manages the creation and setup of the StateGraph based on use cases. | |
| """ | |
| def __init__(self,model): | |
| self.llm = model | |
| self.graph_builder = StateGraph(State) | |
| self.sdlc_graph_builder = StateGraph(SDLCState) | |
| def basic_chatbot_build_graph(self): | |
| """ | |
| Builds a basic chatbot graph using LangGraph. | |
| This method initializes a chatbot node using the `BasicChatbotNode` class | |
| and integrates it into the graph. The chatbot node is set as both the | |
| entry and exit point of the graph. | |
| """ | |
| self.basic_chatbot_node = BasicChatbotNode(self.llm) | |
| self.graph_builder.add_node("chatbot", self.basic_chatbot_node.process) | |
| self.graph_builder.set_entry_point("chatbot") | |
| self.graph_builder.set_finish_point("chatbot") | |
| def chatbot_with_tool_build_graph(self): | |
| """ | |
| Builds an advanced chatbot graph with tool integration. | |
| This method creates a chatbot graph that includes both a chatbot node | |
| and a tool node. It defines tools, initializes the chatbot with tool | |
| capabilities, and sets up conditional and direct edges between nodes. | |
| The chatbot node is set as the entry point. | |
| """ | |
| # Define tools and tool node | |
| tools = get_tools() | |
| tool_node = create_tool_node(tools) | |
| # Define LLM | |
| llm = self.llm | |
| # Define chatbot node | |
| obj_chatbot_with_node = ChatbotWithToolNode(llm) | |
| chatbot_node = obj_chatbot_with_node.create_chatbot(tools) | |
| # Add nodes | |
| self.graph_builder.add_node("chatbot", chatbot_node) | |
| self.graph_builder.add_node("tools", tool_node) | |
| # Define conditional and direct edges | |
| self.graph_builder.add_conditional_edges("chatbot", tools_condition) | |
| self.graph_builder.add_edge("tools", "chatbot") | |
| # Set entry point and compile graph | |
| self.graph_builder.set_entry_point("chatbot") | |
| def travel_planner_build_graph(self): | |
| """ | |
| Builds a standalone travel planning graph with itinerary generation. | |
| """ | |
| # Initialize the Travel Planner node | |
| travel_planner_node = TravelPlannerNode(self.llm) | |
| # Add the Travel Planner node to the graph | |
| self.graph_builder.add_node("travel_planner", travel_planner_node.process) | |
| # Set the entry point to the Travel Planner node | |
| self.graph_builder.set_entry_point("travel_planner") | |
| # Define the edge to end the graph after the Travel Planner completes | |
| self.graph_builder.add_edge("travel_planner", END) | |
| # Helper methods - START | |
| # Nodes | |
| def call_caller_model(self,state: MessagesState): | |
| state["current_time"] = datetime.datetime.now().strftime("%Y-%m-%d %H:%M") | |
| response = self.caller_model.invoke(state) | |
| return {"messages": [response]} | |
| # Edges | |
| def should_continue_caller(self,state: MessagesState): | |
| messages = state["messages"] | |
| last_message = messages[-1] | |
| if not last_message.tool_calls: | |
| return "end" | |
| else: | |
| return "continue" | |
| # Helper method - END | |
| def appointment_receptionist_bot_build_graph(self): | |
| caller_tools = [book_appointment, get_next_available_appointment, cancel_appointment] | |
| tool_node = ToolNode(caller_tools) | |
| caller_pa_prompt = """You are a personal assistant, and need to help the user to book or cancel appointments, you should check the available appointments before booking anything. Be extremely polite, so much so that it is almost rude. | |
| Current time: {current_time} | |
| """ | |
| caller_chat_template = ChatPromptTemplate.from_messages([ | |
| ("system", caller_pa_prompt), | |
| ("placeholder", "{messages}"), | |
| ]) | |
| self.caller_model = caller_chat_template | self.llm.bind_tools(caller_tools) | |
| # Add Nodes | |
| self.graph_builder.add_node("agent", self.call_caller_model) | |
| self.graph_builder.add_node("action", tool_node) | |
| # Add Edges | |
| self.graph_builder.add_conditional_edges( | |
| "agent", | |
| self.should_continue_caller, | |
| { | |
| "continue": "action", | |
| "end": END, | |
| }, | |
| ) | |
| self.graph_builder.add_edge("action", "agent") | |
| # Set Entry Point and build the graph | |
| self.graph_builder.set_entry_point("agent") | |
| def customer_support_build_graph(self): | |
| obj_cs_bot = Customer_Support_Bot(llm=self.llm) | |
| self.graph_builder = obj_cs_bot.chat_bot() | |
| def ai_news_build_graph(self): | |
| # Initialize the AINewsNode | |
| ai_news_node = AINewsNode(self.llm) | |
| self.graph_builder.add_node("fetch_news", ai_news_node.fetch_news) | |
| self.graph_builder.add_node("summarize_news", ai_news_node.summarize_news) | |
| self.graph_builder.add_node("save_result", ai_news_node.save_result) | |
| self.graph_builder.set_entry_point("fetch_news") | |
| self.graph_builder.add_edge("fetch_news", "summarize_news") | |
| self.graph_builder.add_edge("summarize_news", "save_result") | |
| self.graph_builder.add_edge("save_result", END) | |
| def sdlc_workflow_build_graph(self): | |
| sdlc_wf_node = SDLCNode(self.llm) | |
| self.graph_builder = self.sdlc_graph_builder | |
| try: | |
| # Add all primary workflow nodes. | |
| nodes = [ | |
| ("generate_user_stories", sdlc_wf_node.generate_user_stories), | |
| ("product_owner_review", sdlc_wf_node.product_owner_review), | |
| ("create_design_docs", sdlc_wf_node.create_design_docs), | |
| ("revise_user_stories", sdlc_wf_node.revise_user_stories), | |
| ("design_review", sdlc_wf_node.design_review), | |
| ("generate_code", sdlc_wf_node.generate_code), | |
| ("code_review", sdlc_wf_node.code_review), | |
| ("security_review", sdlc_wf_node.security_review), | |
| ("fix_code_after_code_review", sdlc_wf_node.fix_code_after_code_review), | |
| ("fix_code_after_security", sdlc_wf_node.fix_code_after_security), | |
| ("write_test_cases", sdlc_wf_node.write_test_cases), | |
| ("test_cases_review", sdlc_wf_node.test_cases_review), | |
| ("fix_test_cases", sdlc_wf_node.fix_test_cases), | |
| ] | |
| # Helper functions to wrap review nodes. | |
| def human_loop_node(review_field): | |
| def node(state): | |
| # Trigger an interrupt to surface the LLM-generated review. | |
| if st.session_state.user_decision is not "approve": | |
| value = interrupt({ | |
| "__interrupt__": True, | |
| "review": state.get(review_field, ""), | |
| "instruction": f"Please review the '{review_field}'. Approve or provide feedback to reject." | |
| }) | |
| else : | |
| value = st.session_state.user_decision | |
| st.session_state.user_decision = '' | |
| return {"human_decision": value} | |
| return node | |
| def decision_node(previous_node): | |
| def node(state): | |
| if state.get("human_decision") == "approve": | |
| state["decision"] = "approve" | |
| else: | |
| state["decision"] = "reject" | |
| state["feedback"] = state.get("human_decision") | |
| return state | |
| return node | |
| review_nodes = ["product_owner_review", "design_review", "code_review", "test_cases_review"] | |
| additional_nodes = [] | |
| for review in review_nodes: | |
| additional_nodes.append((f"human_loop_{review}", human_loop_node(review))) | |
| # Set the previous node for rejection (adjust as needed): | |
| if review == "product_owner_review": | |
| prev = "generate_user_stories" | |
| elif review == "design_review": | |
| prev = "revise_user_stories" | |
| elif review == "code_review": | |
| prev = "generate_code" | |
| elif review == "test_cases_review": | |
| prev = "write_test_cases" | |
| additional_nodes.append((f"decision_{review}", decision_node(prev))) | |
| # Add all nodes to the graph. | |
| for node_name, node_func in nodes + additional_nodes: | |
| self.graph_builder.add_node(node_name, node_func) | |
| # Set entry point. | |
| if st.session_state.graph_stage == 'resumed': | |
| self.graph_builder.set_entry_point(st.session_state.state['current_step']) | |
| else: | |
| self.graph_builder.set_entry_point("generate_user_stories") | |
| # ---- Build Flow Edges ---- | |
| # Wrap product_owner_review: | |
| self.graph_builder.add_edge("generate_user_stories", "product_owner_review") | |
| self.graph_builder.add_edge("product_owner_review", "human_loop_product_owner_review") | |
| self.graph_builder.add_edge("human_loop_product_owner_review", "decision_product_owner_review") | |
| self.graph_builder.add_conditional_edges( | |
| "decision_product_owner_review", | |
| lambda state: "approve" if state.get("decision") == "approve" else "reject", | |
| { | |
| "approve": "create_design_docs", | |
| "reject": "generate_user_stories" | |
| } | |
| ) | |
| # Wrap design_review: | |
| self.graph_builder.add_edge("revise_user_stories", "design_review") | |
| self.graph_builder.add_edge("design_review", "human_loop_design_review") | |
| self.graph_builder.add_edge("human_loop_design_review", "decision_design_review") | |
| self.graph_builder.add_conditional_edges( | |
| "decision_design_review", | |
| lambda state: "approve" if state.get("decision") == "approve" else "reject", | |
| { | |
| "approve": "generate_code", | |
| "reject": "revise_user_stories" | |
| } | |
| ) | |
| # Wrap code_review: | |
| self.graph_builder.add_edge("generate_code", "code_review") | |
| self.graph_builder.add_edge("code_review", "human_loop_code_review") | |
| self.graph_builder.add_edge("human_loop_code_review", "decision_code_review") | |
| self.graph_builder.add_conditional_edges( | |
| "decision_code_review", | |
| lambda state: "approve" if state.get("decision") == "approve" else "reject", | |
| { | |
| "approve": "security_review", | |
| "reject": "generate_code" | |
| } | |
| ) | |
| # Wrap test_cases_review: | |
| self.graph_builder.add_edge("write_test_cases", "test_cases_review") | |
| self.graph_builder.add_edge("test_cases_review", "human_loop_test_cases_review") | |
| self.graph_builder.add_edge("human_loop_test_cases_review", "decision_test_cases_review") | |
| self.graph_builder.add_conditional_edges( | |
| "decision_test_cases_review", | |
| lambda state: "approve" if state.get("decision") == "approve" else "reject", | |
| { | |
| "approve": "fix_test_cases", | |
| "reject": "write_test_cases" | |
| } | |
| ) | |
| # Other sequential edges. | |
| self.graph_builder.add_edge("create_design_docs", "revise_user_stories") | |
| self.graph_builder.add_edge("security_review", "fix_code_after_code_review") | |
| self.graph_builder.add_edge("fix_code_after_code_review", "fix_code_after_security") | |
| self.graph_builder.add_edge("fix_code_after_security", "write_test_cases") | |
| # Set finish point at the end of the workflow. | |
| self.graph_builder.set_finish_point("fix_test_cases") | |
| except Exception as e: | |
| print(e) | |
| self.graph_builder | |
| def setup_graph(self, usecase: str): | |
| """ | |
| Sets up the graph for the selected use case. | |
| """ | |
| if usecase == "Basic Chatbot": | |
| self.basic_chatbot_build_graph() | |
| elif usecase == "Chatbot with Tool": | |
| self.chatbot_with_tool_build_graph() | |
| elif usecase == "Travel Planner": | |
| self.travel_planner_build_graph() | |
| elif usecase == "Appointment Receptionist": | |
| self.appointment_receptionist_bot_build_graph() | |
| elif usecase =="Customer Support": | |
| self.customer_support_build_graph() | |
| elif usecase =="AI News": | |
| self.ai_news_build_graph() | |
| elif usecase =="SDLC Workflow": | |
| checkpointer = MemorySaver() | |
| self.sdlc_workflow_build_graph() | |
| return self.graph_builder.compile(checkpointer=checkpointer) | |
| else: | |
| raise ValueError("Invalid use case selected.") | |
| return self.graph_builder.compile() | |