Merge pull request #74 from amosproj/API_error_handling_branch

Api error handling branch and context awareness

Author: bergzain

src/ChatUI_streamlit/LLMModel.py

@@ -1,8 +1,6 @@
-
-# LLMModel
+#%%
 import os
 from langchain.chains import RetrievalQA
-
 from langchain.cache import InMemoryCache
 from langchain.globals import set_llm_cache
 from langchain.document_loaders.generic import GenericLoader
@@ -16,89 +14,82 @@
 from langchain.memory import ConversationSummaryMemory
 from langchain.vectorstores import FAISS
 from langchain.document_loaders.text import TextLoader
+from langchain.agents import AgentType, Tool, initialize_agent
+from langchain.memory import ConversationBufferMemory
+from langchain.agents import AgentExecutor
+# import faiss
+from langchain.vectorstores import FAISS as FAISS
+import faiss 
 
-openai_api_key = os.environ["OPENAI_API_KEY"]
+# Load the OpenAI API key
 
-# check if the API key is loaded
+openai_api_key = os.environ["OPENAI_API_KEY"]
 assert openai_api_key is not None, "Failed to load the OpenAI API key from .env file. Please create .env file and add OPENAI_API_KEY = 'your key'"
 
-
-
-
-llm = ChatOpenAI(model_name='gpt-3.5-turbo',openai_api_key=openai_api_key) # Load the LLM model
-# set_llm_cache(InMemoryCache())
-
-
-embeddings = OpenAIEmbeddings(disallowed_special=(), openai_api_key=openai_api_key) # Load the embeddings
-#
-# # This is the root directory for the documents i want to create the RAG from
-# repo_path = '/Users/zainhazzouri/projects/amos2023ws05-pipeline-config-chat-ai/src/RAG'
-# loader = GenericLoader.from_filesystem(
-#     repo_path,
-#     glob="**/*",
-#     suffixes=[".py"],
-#     parser=LanguageParser(language=Language.PYTHON, parser_threshold=500),
-# )
-# documents = loader.load()
-#
-# python_splitter = RecursiveCharacterTextSplitter.from_language(
-#     language=Language.PYTHON, chunk_size=2000, chunk_overlap=200
-# )
-# texts = python_splitter.split_documents(documents)
-#
-#
-# db = Chroma.from_documents(texts, OpenAIEmbeddings(disallowed_special=()))
-# retriever = db.as_retriever(
-#     search_type="mmr",  # Also test "similarity"
-#     search_kwargs={"k": 8},
-# )
-
-##########################################  the old version of RAG
-# This is the root directory for the documents i want to create the RAG from
-root_dir = os.path.join("..", "RAG")
-docs = [] # Create an empty list to store the docs
-
-# Go through each folder to extract all the files
-for dirpath, dirnames, filenames in os.walk(root_dir):
-
-    # Go through each file
-    for file in filenames:
-        try:
-            # Load up the file as a doc and split
-            loader = TextLoader(os.path.join(dirpath, file), encoding='utf-8')
-            docs.extend(loader.load_and_split())
-        except Exception as e:
-            pass
-
-docsearch = FAISS.from_documents(docs, embeddings) # Create the FAISS index
-# source https://python.langchain.com/docs/integrations/vectorstores/faiss_async
-
-
-#memory = ConversationSummaryMemory(llm=llm, memory_key="chat_history", return_messages=True)
-# add caching to the memory
-
-
-RAG = RetrievalQA.from_chain_type(llm,chain_type="stuff" ,retriever=docsearch.as_retriever()) # the old chain for the retrieval
+# Initialize the language model
+llm = ChatOpenAI(model_name='gpt-3.5-turbo', openai_api_key=openai_api_key)
+
+# Load the embeddings
+embeddings = OpenAIEmbeddings(disallowed_special=(), openai_api_key=openai_api_key)
+
+# # Load and split documents
+# root_dir = '/Users/zainhazzouri/projects/amos2023ws05-pipeline-config-chat-ai/src/RAG/pipelines'
+# docs = []
+# for dirpath, dirnames, filenames in os.walk(root_dir):
+#     for file in filenames:
+#         try:
+#             loader = TextLoader(os.path.join(dirpath, file), encoding='utf-8')
+#             docs.extend(loader.load_and_split())
+#         except Exception as e:
+#             pass  # Consider logging the exception for debugging
+
+# # Create the FAISS index
+# docsearch = FAISS.from_documents(docs, embeddings)
+
+#%%
+# save the vector store offline for later use
+# faiss.write_index(docsearch.index, '/Users/zainhazzouri/projects/amos2023ws05-pipeline-config-chat-ai/src/ChatUI_streamlit/faiss_index_file')
+# docsearch.save_local("/Users/zainhazzouri/projects/amos2023ws05-pipeline-config-chat-ai/src/ChatUI_streamlit/faiss_index")
+
+#%%
+docsearch = FAISS.load_local("/Users/zainhazzouri/projects/amos2023ws05-pipeline-config-chat-ai/src/ChatUI_streamlit/faiss_index", embeddings)
+#%%
+# Initialize RetrievalQA
+RAG = RetrievalQA.from_chain_type(llm, chain_type="stuff", retriever=docsearch.as_retriever())
+
+# Define tools
+tools = [
+    Tool(
+        name="RTDIP SDK",
+        func=RAG.run,
+        description="useful for when you need to answer questions about RTDIP",
+    )
+]
+
+# Initialize conversation memory
+conversation_memory = ConversationBufferMemory()
+
+# Initialize Agent with conversation memory
+agent = initialize_agent(
+    tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True, memory=conversation_memory, handle_parsing_errors=True
+)
+
+# Set the LLM cache
 set_llm_cache(InMemoryCache())
 
+# Function to update and retrieve conversation context
+def update_and_get_context(user_input, conversation_memory):
+    conversation_memory.add_user_input(user_input)
+    context = conversation_memory.get_conversation()
+    model_input = "\n".join(context + [user_input])
+    return model_input
 
-#RAG = ConversationalRetrievalChain.from_llm(llm,chain_type="stuff", retriever=docsearch.as_retriever()) # the new chain for the retrieval
-
+# Example usage (commented out for testing)
+# user_input = "What's the weather like today?"
+# model_input = update_and_get_context(user_input, conversation_memory)
+# response = llm.run(model_input)
+# print(response)
 
-##### this code for testing the model don't delete it --
+# Note: You can uncomment and modify the testing code as per your use case.
 
-# question1 = " Hello , my name is Zain"
-# question2 = " what's my name?"
-#question3 = "I would like to use RTDIP components to read from an eventhub using ‘connection string’ as the connection string, and ‘consumer group’ as the consumer group, transform using binary to string, and edge x transformer then write to delta, return only the python code "
-#
-# result = RAG(question1)
-# result["answer"]
-# print(result["answer"])
-#
-# result = RAG(question2)
-# result["answer"]
-# print(result["answer"])
-#
-#result = RAG(question3)
-#result["answer"]
-#print(result["answer"])
+# %%

src/ChatUI_streamlit/app.py

@@ -1,18 +1,29 @@
-
-# app
 import streamlit as st
-import replicate
 import os
 import time
-
-
-
-# App title
+import requests
+import openai
+
+class InvalidAPIKeyException(Exception):
+    pass
+
+# Function to check API key validity
+def is_valid_api_key(key):
+    url = "https://api.openai.com/v1/models/gpt-3.5-turbo-instruct"
+    headers = {"Authorization": f"Bearer {key}"}
+    try:
+        response = requests.get(url, headers=headers)
+        return response.status_code == 200
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return False
+
+# Initialize page configuration once
 if 'page_config_set' not in st.session_state:
     st.set_page_config(page_title="RTDIP Pipeline Chatbot")
     st.session_state['page_config_set'] = True
 
-# Use HTML/CSS to position the title and GitHub link on the same line
+# HTML/CSS for title and GitHub link
 st.markdown(
     '''
     <div style="display: flex; justify-content: space-between; align-items: center;">
@@ -21,24 +32,21 @@
     </div>
     ''', unsafe_allow_html=True)
 
+# Check if the OpenAI API key is already stored in the session
+if 'OPENAI_API_KEY' not in st.session_state:
+    # If not, ask the user to input it
+    openai_api_key = st.text_input('Enter OpenAI API Key:', type='password')
+    if openai_api_key:
+        try:
+            if is_valid_api_key(openai_api_key):
+                st.session_state['OPENAI_API_KEY'] = openai_api_key
+                os.environ['OPENAI_API_KEY'] = openai_api_key
+                st.success('API Key stored!')
+            else:
+                raise InvalidAPIKeyException
+        except InvalidAPIKeyException:
+            st.error('Invalid OpenAI API Key. Please enter a valid key.')
 
-# Replicate Credentials
-api_key_container = st.empty()
-openai_api_key = api_key_container.text_input('Enter OpenAI API Key:', type='password')
-
-# Check if OpenAI API Key is entered
-if openai_api_key:
-        # Store the API key in the session state 
-        st.session_state['OPENAI_API_KEY'] = openai_api_key
-        os.environ['OPENAI_API_KEY'] = openai_api_key
-        success_message = st.success('API Key stored!')
-        # Hide success message, input field, and chat messages after 3 seconds
-        time.sleep(0)
-        success_message.empty()
-        api_key_container.empty()
-else:
-        st.warning('Invalid OpenAI API Key. Please enter a valid key.')
-        
 # Store LLM generated responses
 if "conversations" not in st.session_state.keys():
     st.session_state.conversations = [{"title": "Default Conversation", "messages": [{"role": "assistant", "content": "How may I assist you today?"}]}]
@@ -49,48 +57,27 @@
         with st.chat_message(message["role"]):
             st.write(message["content"])
 
-def clear_chat_history():
-    st.session_state.conversations = [{"title": "Default Conversation", "messages": [{"role": "assistant", "content": "How may I assist you today?"}]}]
-#st.sidebar.button('Clear Chat History', on_click=clear_chat_history)
-
-
 # User-provided prompt
 if 'OPENAI_API_KEY' in st.session_state and st.session_state['OPENAI_API_KEY']:
     from LLMModel import RAG as RAG
     if prompt := st.chat_input():
-    # Get the conversation context
         conversation = st.session_state.conversations[-1]
-    
-    # Use the entire conversation context as input
-        #context = "\n".join([message["content"] for message in conversation["messages"]])
-    
-    # Add the user's input to the conversation
+        context = "\n".join([message["content"] for message in conversation["messages"]])
         conversation["messages"].append({"role": "user", "content": prompt})
-
-    # Display user's input in the chat
         with st.chat_message("user"):
             st.write(prompt)
-
-    # Generate a new response considering the entire conversation context
         with st.chat_message("assistant"):
-            start_time = time.time()  # to calculate the time taken to generate the response
+            start_time = time.time()
             with st.spinner("Generating..."):
-                response = RAG.run(prompt)
-                end_time = time.time()  # to calculate the time taken to generate the response
+                response = RAG.run(context + "\n" + prompt)
+                end_time = time.time()
                 placeholder = st.empty()
                 full_response = ''
                 for item in response:
                     full_response += item
                     placeholder.markdown(full_response)
                 placeholder.markdown(full_response)
-
-    # Calculate the time taken
         response_time = end_time - start_time
         st.write(f"Response generated in {response_time:.2f} seconds.")
-
-    # Add the assistant's response to the conversation
         message = {"role": "assistant", "content": full_response}
         conversation["messages"].append(message)
-    
-    
-