Upload 10 files

ensemble.py

@@ -1,4 +1,5 @@
 import os
+import logging
 
 from langchain_community.retrievers import BM25Retriever, TavilySearchAPIRetriever
 from langchain.retrievers import EnsembleRetriever
@@ -20,7 +21,7 @@ def ensemble_retriever_from_docs(docs, embeddings=None):
     bm25_retriever = BM25Retriever.from_texts([t.page_content for t in texts])
 
     # tavily_retriever = TavilySearchAPIRetriever(k=3, include_domains=['https://ilibrary.ru/text/107'])
-    tavily_retriever = TavilySearchAPIRetriever(k=3, include_domains=['https://equitygroupholdings.com'])
+    tavily_retriever = MyTavilySearchAPIRetriever(k=3, include_domains=['https://equitygroupholdings.com'])
 
     ensemble_retriever = EnsembleRetriever(
         retrievers=[bm25_retriever, vs_retriever, tavily_retriever],
@@ -29,6 +30,17 @@ def ensemble_retriever_from_docs(docs, embeddings=None):
     return ensemble_retriever
 
 
+class MyTavilySearchAPIRetriever(TavilySearchAPIRetriever):
+    def _get_relevant_documents(
+        self, query: str, *, run_manager
+    ):
+        try:
+            return super()._get_relevant_documents(query, run_manager=run_manager)
+        except Exception as e:
+            logging.error(f"TavilySearch error: {e}")
+            return []
+
+
 def main():
     load_dotenv()
 

full_chain.py

@@ -44,7 +44,7 @@ def create_full_chain(retriever, openai_api_key=None):
 
 def ask_question(chain, query, session_id):
     # try:
-    logging.info(f"Send request: {query}")
+    # logging.info(f"Send request from session {session_id}: {query}")
     response = chain.invoke(
         {"question": query},
         config={"configurable": {"session_id": session_id}}

gradio_app.py

@@ -0,0 +1,216 @@
+import os
+import streamlit as st
+import logging
+from langchain_community.chat_message_histories import StreamlitChatMessageHistory
+from langchain_community.embeddings import OpenAIEmbeddings, HuggingFaceEmbeddings
+from langchain_community.retrievers import BM25Retriever
+
+from ensemble import ensemble_retriever_from_docs
+from full_chain import create_full_chain, ask_question
+from local_loader import load_data_files, load_file
+from vector_store import EmbeddingProxy 
+from memory import clean_session_history
+from pathlib import Path
+
+import gradio as gr
+from langchain.chat_models import ChatOpenAI
+from langchain.schema import AIMessage, HumanMessage
+
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+def show_ui(message, history, request: gr.Request):
+    """
+    Displays the Streamlit chat UI and handles user interactions.
+
+    Args:
+        qa: The LangChain chain for question answering.
+        prompt_to_user: The initial prompt to display to the user.
+    """
+    global chain
+    session_id = request.session_hash
+    response = ask_question(chain, message, session_id)
+    # logging.info(f"Response: {response}")
+    return response.content
+
+
+def get_retriever(openai_api_key=None):
+    """
+    Creates and caches the document retriever.
+
+    Args:
+        openai_api_key: The OpenAI API key.
+
+    Returns:
+        An ensemble document retriever.
+    """
+    try:
+        docs = load_data_files(data_dir="data")  
+        # embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key, model="text-embedding-3-small")
+        embeddings = HuggingFaceEmbeddings()
+        return ensemble_retriever_from_docs(docs, embeddings=embeddings)
+    except Exception as e:
+        logging.error(f"Error creating retriever: {e}")
+        logging.exception(f"message")
+        st.error("Error initializing the application. Please check the logs.")
+        st.stop()  # Stop execution if retriever creation fails
+
+
+def get_chain(openai_api_key=None, huggingfacehub_api_token=None):
+    """
+    Creates the question answering chain.
+
+    Args:
+        openai_api_key: The OpenAI API key.
+        huggingfacehub_api_token: The Hugging Face Hub API token.
+
+    Returns:
+        A LangChain question answering chain.
+    """
+    try:
+        ensemble_retriever = get_retriever(openai_api_key=openai_api_key)
+        chain = create_full_chain(
+            ensemble_retriever,
+            openai_api_key=openai_api_key,
+        )
+        return ensemble_retriever, chain
+    except Exception as e:
+        logging.error(f"Error creating chain: {e}")
+        logging.exception(f"message")
+        st.error("Error initializing the application. Please check the logs.")
+        st.stop()  # Stop execution if chain creation fails
+
+def get_secret_or_input(secret_key, secret_name, info_link=None):
+    """
+    Retrieves a secret from Streamlit secrets or prompts the user for input.
+
+    Args:
+        secret_key: The key of the secret in Streamlit secrets.
+        secret_name: The user-friendly name of the secret.
+        info_link: An optional link to provide information about the secret.
+
+    Returns:
+        The secret value.
+    """
+    if secret_key in st.secrets:
+        st.write("Found %s secret" % secret_key)
+        secret_value = st.secrets[secret_key]
+    else:
+        st.write(f"Please provide your {secret_name}")
+        secret_value = st.text_input(secret_name, key=f"input_{secret_key}", type="password")
+        if secret_value:
+            st.session_state[secret_key] = secret_value
+        if info_link:
+            st.markdown(f"[Get an {secret_name}]({info_link})")
+    return secret_value
+
+def process_uploaded_file(uploaded_file):
+    """
+    Processes the uploaded file and adds it to the vector database.
+
+    Args:
+        uploaded_file: The uploaded file object from Streamlit.
+        openai_api_key: The OpenAI API key for embedding generation.
+    """
+    # try:
+    if uploaded_file is not None:
+        logging.info(f'run upload {uploaded_file}')
+
+        if isinstance(uploaded_file, str):
+            filename = uploaded_file
+        else:
+            filename = str(uploaded_file.name)
+
+        # Load the document using the saved file path
+        docs = load_file(Path(filename))
+
+        global ensemble_retriever
+        global chain
+
+        all_docs = ensemble_retriever.retrievers[0].docs
+        all_docs.extend(docs)
+
+        ensemble_retriever.retrievers[1].add_documents(docs)
+
+        new_bm25 = BM25Retriever.from_texts([t.page_content for t in all_docs])
+
+        ensemble_retriever.retrievers[0] = new_bm25
+
+        chain = create_full_chain(
+            ensemble_retriever,
+            openai_api_key=open_api_key,
+        )
+
+        logging.info("File uploaded and added to the knowledge base!")
+        gr.Info('File uploaded and added to the knowledge base!', duration=3)
+    
+    return None
+        
+    # except Exception as e:
+    #     logging.error(f"Error processing uploaded file: {e}")
+    #     st.error("Error processing the file. Please check the logs.")
+
+SUPPORTED_FORMATS = ['.txt', '.json', '.pdf']
+
+def activate():
+    return gr.update(interactive=True)
+
+def deactivate():
+    return gr.update(interactive=False)
+
+def reset(z, request: gr.Request):
+    session_id = request.session_hash
+    clean_session_history(session_id)
+    return [], []
+
+def main():
+    with gr.Blocks() as demo:
+        gr.Markdown(
+            "# Equity Bank AI assistant \n"
+            "Ask questions about Equity Bank's products and services:"
+        )
+        with gr.Tab('Chat'):
+            clean_btn = gr.Button(value="Clean history", variant="secondary", size='sm', render=False)
+            bot = gr.Chatbot(elem_id="chatbot", render=False)
+
+            chat = gr.ChatInterface(
+                show_ui,
+                chatbot=bot,
+                undo_btn=None,
+                retry_btn=None,
+                clear_btn=clean_btn,
+            )
+        with gr.Tab('Documents'):
+            file_input = gr.File(
+                label=f'{", ".join([str(f) for f in SUPPORTED_FORMATS])}',
+                file_types=SUPPORTED_FORMATS,
+            )
+            submit_btn = gr.Button(value="Index file", variant="primary", interactive=False)
+
+        clean_btn.click(fn=reset, inputs=clean_btn, outputs=[bot, chat.chatbot_state])
+
+        submit_btn.click(
+                fn=process_uploaded_file,
+                inputs=file_input,
+                outputs=file_input,
+                api_name="Index file"
+            )
+        
+        file_input.upload(fn=activate, outputs=[submit_btn])
+        file_input.clear(fn=deactivate, outputs=[submit_btn])
+
+    demo.launch(share=True)
+
+
+open_api_key = os.getenv('OPEN_API_KEY')
+
+ensemble_retriever, chain = get_chain(
+    openai_api_key=open_api_key,
+    huggingfacehub_api_token=None
+)
+
+
+
+if __name__ == "__main__":
+    main()

memory.py

@@ -39,8 +39,6 @@ def create_memory_chain(llm, base_chain):
         if session_id not in store:
             store[session_id] = ChatMessageHistory()
 
-        logging.info(str(store))
-
         return store[session_id]
 
     with_message_history = RunnableWithMessageHistory(

requirements.txt

@@ -1,11 +1,11 @@
-chromadb
-huggingface-hub
-langchain
-langchain-community
-langchain-openai
-sentence-transformers
-streamlit
-gradio
-pypdf
-rank_bm25
-tavily-python
+chromadb==0.5.5
+huggingface-hub==0.24.6
+langchain==0.2.14
+langchain-community==0.2.12
+langchain-openai==0.1.22
+sentence-transformers==3.0.1
+streamlit==1.37.1
+gradio==4.41.0
+pypdf==4.3.1
+rank_bm25==0.2.2
+tavily-python==0.4.0