Update app.py

app.py

@@ -1,165 +1,168 @@
-import os
-import streamlit as st
-import openai
-from openai import OpenAI
-from langchain_openai import OpenAIEmbeddings
-from langchain_openai import ChatOpenAI
-
-from pinecone import Pinecone
-from langchain_pinecone import PineconeVectorStore
-
-from agents import Obnoxious_Agent, Query_Agent, Relevant_Documents_Agent, Answering_Agent
-
-client = None
-
-# **Streamlit UI Elements**
-st.title("Mini Project 2: Streamlit Chatbot")
-
-# Pinecone Setup**
-pc = Pinecone(api_key=st.secrets['PINECONE_API_KEY'])
-index = pc.Index('openaiembedding')
-
-# Session State for API key and Chat history
-if "is_valid" not in st.session_state:
-    st.session_state["is_valid"] = False
-if 'api_key' not in st.session_state:
-    st.session_state["api_key"] = ""
-if "messages" not in st.session_state:
-    st.session_state['messages'] = []
-
-def check_openai_api_key():
-    client = OpenAI(api_key=st.session_state["api_key"])
-    st.session_state["is_valid"] = False
-    try:
-        client.models.list()
-    except openai.APIError as e:
-        st.session_state["is_valid"] = False
-    except openai.APIConnectionError as e:
-        st.session_state["is_valid"] = False
-    except openai.AuthenticationError as e:
-        st.session_state["is_valid"] = False
-    else:
-        st.session_state["is_valid"] = True
-        st.success("Valid OpenAI API key entered successfully!")
-
-if not st.session_state["is_valid"]:
-    st.session_state["api_key"] = st.text_input(label="Enter your OpenAI API Key", type="password")
-    check_openai_api_key() 
-
-class Head_Agent:
-    def __init__(self, client, index, embeddings, api_key) -> None:
-        # TODO: Initialize the Head_Agent
-        # OpenAI client + embeddings
-        self.client = client
-        self.model = ChatOpenAI(api_key=api_key)
-        self.embeddings = embeddings
-        # Pinecone Index + vector store
-        self.index =  index
-        self.vectorstore = PineconeVectorStore(index_name='openaiembedding', embedding=self.embeddings, namespace="embedding_2048_256", text_key="text")
-        self.retriever = self.vectorstore.as_retriever()
-        # Filter user queries
-        self.obnoxious_phrase = "Please do not ask inappropriate or obnoxious questions."
-        self.non_relevant_phrase = "No relevant documents found in the documents. Please ask a question relevant to the book on Machine Learning"
-        # chat history prompt
-        self.memory_prompt = """Given a chat history and the latest user question \
-        which might reference context in the chat history, formulate a standalone question \
-        which can be understood without the chat history. Do NOT answer the question, \
-        just reformulate it if needed and otherwise return it as is.
-        Chat_History: {chat_history}
-        """
-
-    def setup_sub_agents(self):
-        self.OA = Obnoxious_Agent(self.client)
-        self.QA = Query_Agent(self.client, self.embeddings, self.index, self.vectorstore)
-        self.RDA = Relevant_Documents_Agent(self.client, self.embeddings)
-        self.AA = Answering_Agent(self.model, self.retriever)
-        print("Agents Initialized")
-
-    def query_gpt(self, prompt) -> str:
-        # TODO: Get if the returned documents are relevant
-        try:
-            response = self.client.chat.completions.create(model="gpt-3.5-turbo",
-                                                                  temperature=0,
-                                                                messages=[{
-                                                                            "role": "system",
-                                                                            "content": prompt
-                                                                          }]
-                                                                 )
-            return response.choices[0].message.content
-        except Exception as e:
-            print(f"Error while checking with gpt: {e}")
-            return None
-        
-    def set_chat_history(self, chat_history):
-        self.memory_prompt = self.memory_prompt.format(chat_history=chat_history)
-
-    def main_loop(self, query, chat_history):
-        # TODO: Get Query from Streamlit App
-        # print("Query:", query)
-
-        # Query for greeting
-        if query.lower() in ["hello", "hi", "hey"]:
-            return "Hello! How can I help you today?"
-
-        ## 1. Obnoxious Agent
-        # Return yes/True = end
-        # Return no/False = continue
-        OA_prompt = self.OA.set_prompt(query)
-        # print("OA Prompt:", OA_prompt)
-        OA_response = self.query_gpt(OA_prompt)
-        print("Obnoxious Response:", OA_response)
-        OA_action = self.OA.extract_action(OA_response)
-        # print("OA Action:", OA_action)
-        if OA_action:
-            print("OA Message:", self.obnoxious_phrase)
-            return self.obnoxious_phrase
-        
-        self.set_chat_history(chat_history)
-        updated_prompt = self.query_gpt(self.memory_prompt)
-
-        ## 2. Pinecone Query - Get Documents
-        pinecone_docs_response = self.QA.query_vector_store(updated_prompt)
-        # print("Documents:", pinecone_docs_response)
-
-        ## 3. Relevant
-        # If relevant continue, else
-        RDA_prompt = self.RDA.set_prompt(pinecone_docs_response, updated_prompt)
-        # print("RDA Prompt:", RDA_prompt)
-        relevance_response = self.query_gpt(RDA_prompt)
-        print("Relevance Response:", relevance_response)
-        RDA_action = self.RDA.extract_action(relevance_response)
-        # print("RDA Action:", RDA_action)
-        if not RDA_action:
-            print("RDA Message:", self.non_relevant_phrase)
-            return self.non_relevant_phrase
-
-        ## 4. Answering
-        AA_prompt = self.AA.set_prompt(updated_prompt)
-        # print("AA Prompt:", AA_prompt)
-        AA_response = self.AA.generate_response()
-        # print("AA Response:", AA_response)
-        return AA_response
-    
-# Initialize Head_Agent if the API key is valid
-if st.session_state["is_valid"]:
-    api_key = st.session_state["api_key"]
-    client = OpenAI(api_key=api_key)
-    embeddings = OpenAIEmbeddings(api_key=api_key, model='text-embedding-ada-002')
-    
-    HA = Head_Agent(client, index, embeddings, api_key)
-    HA.setup_sub_agents()
-
-    # Display existing chat messages
-    for message in st.session_state['messages']:
-        st.chat_message(message['role']).write(message['content'])
-
-    # Main chat interaction loop
-    if user_prompt := st.chat_input("What would you like to ask?"):
-        st.session_state['messages'].append({"role": "user", "content": user_prompt})
-        st.chat_message(st.session_state['messages'][-1]['role']).write(st.session_state['messages'][-1]['content'])
-
-        print("Chat History:", st.session_state['messages'])
-        response = HA.main_loop(query=user_prompt, chat_history=st.session_state['messages'])
-
-        st.session_state['messages'].append({"role": "assistant", "content": response})
-        st.chat_message("assistant").write(response)
+import os
+import streamlit as st
+import openai
+from openai import OpenAI
+from langchain_openai import OpenAIEmbeddings
+from langchain_openai import ChatOpenAI
+
+from pinecone import Pinecone
+from langchain_pinecone import PineconeVectorStore
+
+from agents import Obnoxious_Agent, Query_Agent, Relevant_Documents_Agent, Answering_Agent
+
+client = None
+
+# **Streamlit UI Elements**
+st.title("Mini Project 2: Streamlit Chatbot")
+
+# Pinecone Setup**
+pc = Pinecone(api_key=st.secrets['PINECONE_API_KEY'])
+index = pc.Index('openaiembedding')
+
+# Session State for API key and Chat history
+if "is_valid" not in st.session_state:
+    st.session_state["is_valid"] = False
+if 'api_key' not in st.session_state:
+    st.session_state["api_key"] = ""
+if "messages" not in st.session_state:
+    st.session_state['messages'] = []
+
+def check_openai_api_key():
+    client = OpenAI(api_key=st.session_state["api_key"])
+    st.session_state["is_valid"] = False
+    try:
+        client.models.list()
+    except openai.APIError as e:
+        st.session_state["is_valid"] = False
+    except openai.APIConnectionError as e:
+        st.session_state["is_valid"] = False
+    except openai.AuthenticationError as e:
+        st.session_state["is_valid"] = False
+    else:
+        st.session_state["is_valid"] = True
+        st.success("Valid OpenAI API key entered successfully!")
+
+if not st.session_state["is_valid"]:
+    st.session_state["api_key"] = st.text_input(label="Enter your OpenAI API Key", type="password")
+    check_openai_api_key() 
+
+class Head_Agent:
+    def __init__(self, client, index, embeddings, api_key) -> None:
+        # TODO: Initialize the Head_Agent
+        # OpenAI client + embeddings
+        self.client = client
+        self.model = ChatOpenAI(api_key=api_key)
+        self.embeddings = embeddings
+        # Pinecone Index + vector store
+        self.index =  index
+        self.vectorstore = PineconeVectorStore(index_name='openaiembedding', embedding=self.embeddings, namespace="embedding_2048_256", text_key="text")
+        self.retriever = self.vectorstore.as_retriever()
+        # Filter user queries
+        self.obnoxious_phrase = "Please do not ask inappropriate or obnoxious questions."
+        self.non_relevant_phrase = "No relevant documents found in the documents. Please ask a question relevant to the book on Machine Learning"
+        # chat history prompt
+        self.memory_prompt = """Given a chat history and the latest user question \
+        which might reference context in the chat history, formulate a standalone question \
+        which can be understood without the chat history. Do NOT answer the question, \
+        just reformulate it if needed and otherwise return it as is.
+        Chat_History: {chat_history}
+        """
+
+    def setup_sub_agents(self):
+        self.OA = Obnoxious_Agent(self.client)
+        self.QA = Query_Agent(self.client, self.embeddings, self.index, self.vectorstore)
+        self.RDA = Relevant_Documents_Agent(self.client, self.embeddings)
+        self.AA = Answering_Agent(self.model, self.retriever)
+        print("Agents Initialized")
+
+    def query_gpt(self, prompt) -> str:
+        # TODO: Get if the returned documents are relevant
+        try:
+            response = self.client.chat.completions.create(model="gpt-3.5-turbo",
+                                                                  temperature=0,
+                                                                messages=[{
+                                                                            "role": "system",
+                                                                            "content": prompt
+                                                                          }]
+                                                                 )
+            resp = response.choices[0].message.content
+            if resp is None:
+                resp = "Return none API response empty - possible latency issues."
+            return resp
+        except Exception as e:
+            print(f"Error while checking with gpt: {e}")
+            return None
+        
+    def set_chat_history(self, chat_history):
+        self.memory_prompt = self.memory_prompt.format(chat_history=chat_history)
+
+    def main_loop(self, query, chat_history):
+        # TODO: Get Query from Streamlit App
+        # print("Query:", query)
+
+        # Query for greeting
+        if query.lower() in ["hello", "hi", "hey"]:
+            return "Hello! How can I help you today?"
+
+        ## 1. Obnoxious Agent
+        # Return yes/True = end
+        # Return no/False = continue
+        OA_prompt = self.OA.set_prompt(query)
+        # print("OA Prompt:", OA_prompt)
+        OA_response = self.query_gpt(OA_prompt)
+        print("Obnoxious Response:", OA_response)
+        OA_action = self.OA.extract_action(OA_response)
+        # print("OA Action:", OA_action)
+        if OA_action:
+            print("OA Message:", self.obnoxious_phrase)
+            return self.obnoxious_phrase
+        
+        self.set_chat_history(chat_history)
+        updated_prompt = self.query_gpt(self.memory_prompt)
+
+        ## 2. Pinecone Query - Get Documents
+        pinecone_docs_response = self.QA.query_vector_store(updated_prompt)
+        # print("Documents:", pinecone_docs_response)
+
+        ## 3. Relevant
+        # If relevant continue, else
+        RDA_prompt = self.RDA.set_prompt(pinecone_docs_response, updated_prompt)
+        # print("RDA Prompt:", RDA_prompt)
+        relevance_response = self.query_gpt(RDA_prompt)
+        print("Relevance Response:", relevance_response)
+        RDA_action = self.RDA.extract_action(relevance_response)
+        # print("RDA Action:", RDA_action)
+        if not RDA_action:
+            print("RDA Message:", self.non_relevant_phrase)
+            return self.non_relevant_phrase
+
+        ## 4. Answering
+        AA_prompt = self.AA.set_prompt(updated_prompt)
+        # print("AA Prompt:", AA_prompt)
+        AA_response = self.AA.generate_response()
+        # print("AA Response:", AA_response)
+        return AA_response
+    
+# Initialize Head_Agent if the API key is valid
+if st.session_state["is_valid"]:
+    api_key = st.session_state["api_key"]
+    client = OpenAI(api_key=api_key)
+    embeddings = OpenAIEmbeddings(api_key=api_key, model='text-embedding-ada-002')
+    
+    HA = Head_Agent(client, index, embeddings, api_key)
+    HA.setup_sub_agents()
+
+    # Display existing chat messages
+    for message in st.session_state['messages']:
+        st.chat_message(message['role']).write(message['content'])
+
+    # Main chat interaction loop
+    if user_prompt := st.chat_input("What would you like to ask?"):
+        st.session_state['messages'].append({"role": "user", "content": user_prompt})
+        st.chat_message(st.session_state['messages'][-1]['role']).write(st.session_state['messages'][-1]['content'])
+
+        print("Chat History:", st.session_state['messages'])
+        response = HA.main_loop(query=user_prompt, chat_history=st.session_state['messages'])
+
+        st.session_state['messages'].append({"role": "assistant", "content": response})
+        st.chat_message("assistant").write(response)