|
|
from langgraph.graph import START, END, StateGraph |
|
|
from langgraph.prebuilt import ToolNode, tools_condition |
|
|
from langchain_core.prompts import ChatPromptTemplate |
|
|
import datetime |
|
|
from src.state.state import State |
|
|
from src.nodes.basic_chatbot import BasicChatbot |
|
|
from src.nodes.websearch_chatbot import WebSearchChatbot |
|
|
|
|
|
|
|
|
class GraphBuilder: |
|
|
|
|
|
"""Class to build the state graph for the application.""" |
|
|
|
|
|
def __init__(self, model, session_id: str = "default", tavily_api_key: str = None): |
|
|
self.llm = model |
|
|
self.session_id = session_id |
|
|
self.tavily_api_key = tavily_api_key |
|
|
self.graph_builder = StateGraph(State) |
|
|
|
|
|
def basic_chatbot(self): |
|
|
"""Initialize the basic chatbot node in the graph.""" |
|
|
self.basic_chatbot_node = BasicChatbot(self.llm, self.session_id) |
|
|
self.graph_builder.add_node('basic_chatbot', self.basic_chatbot_node.process) |
|
|
self.graph_builder.add_edge(START, 'basic_chatbot') |
|
|
self.graph_builder.add_edge('basic_chatbot', END) |
|
|
|
|
|
def websearch_chatbot(self): |
|
|
self.websearch_chatbot_node = WebSearchChatbot(self.llm, self.session_id, self.tavily_api_key) |
|
|
self.graph_builder.add_node('websearch_chatbot', self.websearch_chatbot_node.process) |
|
|
self.graph_builder.add_edge(START, 'websearch_chatbot') |
|
|
self.graph_builder.add_edge('websearch_chatbot', END) |
|
|
|
|
|
def setup_graph(self, use_case: str): |
|
|
""" |
|
|
Setup the graph with the appropriate nodes based on use case. |
|
|
|
|
|
:param use_case: The use case for which the graph is being built. |
|
|
""" |
|
|
|
|
|
if use_case == 'Basic Chatbot': |
|
|
self.basic_chatbot() |
|
|
elif use_case == 'Chatbot with Web Search': |
|
|
self.websearch_chatbot() |
|
|
else: |
|
|
self.basic_chatbot() |
|
|
|
|
|
|
|
|
return self.graph_builder.compile() |
