works

src/graph/graph_builder.py

@@ -0,0 +1,38 @@
+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
+
+
+class GraphBuilder:
+    
+    """Class to build the state graph for the application."""
+
+    def __init__(self, model):
+        self.llm = model
+        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.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 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()
+        else:
+            # Default to basic chatbot if use case is not recognized
+            self.basic_chatbot()
+        
+        # Compile and return the graph
+        return self.graph_builder.compile()

src/main.py

@@ -3,7 +3,10 @@ import json
 
 from src.ui.load import LoadStreamlitUI
 from src.llms.groq import GroqLLM
-
+from src.llms.openai import OpenAILLM
+from src.graph.graph_builder import GraphBuilder
+from src.ui.display_results import DisplayResults
+from langchain_core.messages import HumanMessage
 
 def load_app():
     """
@@ -17,26 +20,74 @@ def load_app():
         st.error("Failed to load the UI. Please check your configuration.")
         return
     
-    if st.session_state.IsFetchButtonClicked:
-        user_message = st.session_state.timeframe
+    # Initialize chat history in session state
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    
+    # Display chat history
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.write(message["content"])
     
-    else:
-        user_message = st.text_input("Enter your message:")
+    # Chat input at the bottom
+    if prompt := st.chat_input("Enter your message..."):
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        
+        # Display user message
+        with st.chat_message("user"):
+            st.write(prompt)
         
-    if user_message:
+        # Process the message
+        user_message = prompt
         try:
-            llm = GroqLLM(user_controls_input=user_input)
-            groq_model = llm.get_llm_model()
+            # Get the selected LLM provider from user input
+            selected_llm = user_input.get('Selected LLM')
+            
+            # Initialize the appropriate LLM class based on user selection
+            if selected_llm == 'Groq':
+                llm = GroqLLM(user_controls_input=user_input)
+            elif selected_llm == 'OpenAI':
+                llm = OpenAILLM(user_controls_input=user_input)
+            else:
+                st.error("Please select a valid LLM provider (Groq or OpenAI).")
+                return
             
-            if not groq_model:
-                st.error("Failed to initialize the Groq model. Please check your API key and model selection.")
+            # Get the model instance
+            model = llm.get_llm_model()
+            
+            if not model:
+                st.error(f"Failed to initialize the {selected_llm} model. Please check your API key and model selection.")
                 return
             
-            use_case = user_input.get('use_case')
+            use_case = user_input.get('Selected Use Case')
 
-            if use_case:
-                st.error("Error: No usecase selected.")
+            if not use_case:
+                st.error("Error: No use case selected.")
                 return 
-        
+          
+            graph_builder = GraphBuilder(model=model)
+
+            try:
+                graph = graph_builder.setup_graph(use_case=use_case)
+                
+                # Process the message through the graph
+                ai_response = ""
+                for event in graph.stream({'messages': [HumanMessage(content=user_message)]}):
+                    for value in event.values():
+                        ai_response = value['messages'].content
+                
+                # Add AI response to chat history
+                st.session_state.messages.append({"role": "assistant", "content": ai_response})
+                
+                # Display AI response
+                with st.chat_message("assistant"):
+                    st.write(ai_response)
+                    
+            except Exception as e:
+                st.error(f"An error occurred while setting up the graph: {e}")
+                return
+
         except Exception as e:
-            raise ValueError(f"An error occurred while processing the input: {e}")
+            st.error(f"An error occurred while processing the input: {e}")
+            return

src/nodes/basic_chatbot.py

@@ -0,0 +1,25 @@
+from src.state.state import State
+
+class BasicChatbot:
+    """
+    Class to handle the basic chatbot functionality.
+    """
+
+    def __init__(self, model):
+        """
+        Initialize the BasicChatbot with the given model.
+        
+        :param model: The LLM to be used for the chatbot.
+        """
+        self.model = model
+
+    def process(self, state):
+        """
+        Process the state to generate a response from the model.
+        
+        :param state: The current state of the chatbot.
+        :return: The response generated by the model.
+        """
+        
+        return {'messages': self.model.invoke(state['messages'])}
+       

src/state/state.py

@@ -0,0 +1,9 @@
+from typing import Annotated
+from typing_extensions import TypedDict
+from langgraph.graph.message import add_messages
+
+class State(TypedDict):
+    """
+    State for the graph.
+    """
+    messages: Annotated[list, add_messages]  # List of messages in the graph

src/ui/display_results.py

@@ -0,0 +1,40 @@
+import streamlit as st
+
+from langchain_core.messages import AIMessage, HumanMessage
+
+
+class DisplayResults:  
+  """Class to display results in the Streamlit UI.""" 
+  def __init__(self, user_message, use_case, graph):
+    """
+    Initialize the DisplayResults with user message, use case, and graph.
+    
+    :param user_message: The message input by the user.
+    :param use_case: The selected use case for the chatbot.
+    :param graph: The state graph to be processed.
+    """
+    self.user_message = user_message
+    self.use_case = use_case
+    self.graph = graph
+
+  def display_results(self):
+    """
+    Display the results in the Streamlit UI.
+    
+    This method processes the user message through the graph and displays the response.
+    """
+    use_case = self.use_case
+    user_message = self.user_message
+    graph = self.graph
+
+    if use_case == 'Basic Chatbot':
+        for event in graph.stream({'messages':("user",user_message)}):
+                    print(event.values())
+                    for value in event.values():
+                        print(value['messages'])
+                        with st.chat_message("user"):
+                            st.write(user_message)
+                        with st.chat_message("assistant"):
+                            st.write(value["messages"].content)
+
+