statefulChatbot / src /graph /graph_builder.py
shubhamchau222
chabot graph build & tested
58617a7
import os
# Add the src directory to the Python path
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
from langgraph.graph import StateGraph, START,END
from langgraph.prebuilt import tools_condition,ToolNode
from langchain_core.prompts import ChatPromptTemplate
from src.state.state import graphstate
from src.nodes.basic_chatbot_nodes import BasicChatbotNodes
from src.nodes.chabot_with_tool_nodes import ChatbotWithToolNodes
from src.tools.search_tool import get_tools, create_tool_node
class GraphBuilder:
def __init__(self, model):
self.model = model
self.graph_builder =StateGraph(graphstate)
def basic_chabot_build(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_chabot_node = BasicChatbotNodes(self.model)
self.graph_builder.add_node("chatbot", self.basic_chabot_node.process)
self.graph_builder.add_edge(START, "chatbot")
self.graph_builder.add_edge("chatbot", END)
def chatbot_with_tool_build(self):
"""
Builds a chatbot graph with a tool node using LangGraph.
This method initializes a chatbot node using the `ChatbotWithToolNodes`"
"""
tools:list= get_tools()
tool_node = create_tool_node(tools)
llm= self.model
# Define chatbot node
obj_chatbot_with_node = ChatbotWithToolNodes(llm)
chatbot_node = obj_chatbot_with_node.create_chatbot(tools)
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_edge(START,"chatbot")
self.graph_builder.add_conditional_edges("chatbot", tools_condition)
self.graph_builder.add_edge("tools","chatbot")
def setup_graph(self, usecase: str):
"""
Sets up the graph for the selected use case.
"""
if usecase == "Basic Chatbot":
self.basic_chabot_build()
if usecase == "Chatbot with Tool":
self.chatbot_with_tool_build()
return self.graph_builder.compile()