import os import sys import yaml import gradio as gr import uuid current_dir = os.path.dirname(os.path.abspath(__file__)) from src.document_retrieval import DocumentRetrieval from utils.parsing.sambaparse import parse_doc_universal # added from utils.vectordb.vector_db import VectorDb def handle_userinput(user_question, conversation_chain, history): if user_question: try: # Generate response response = conversation_chain.invoke({"question": user_question}) # Append user message and response to chat history history = history + [(user_question, response["answer"])] return history, "" except Exception as e: error_msg = f"An error occurred: {str(e)}" history = history + [(user_question, error_msg)] return history, "" else: return history, "" def process_documents(files, collection_name, document_retrieval, vectorstore, conversation_chain, api_key=None): try: if api_key: sambanova_api_key = api_key else: sambanova_api_key = os.environ.get('SAMBANOVA_API_KEY') document_retrieval = DocumentRetrieval(sambanova_api_key) _, _, text_chunks = parse_doc_universal(doc=files) print(f'nb of chunks: {len(text_chunks)}') embeddings = document_retrieval.load_embedding_model() collection_id = str(uuid.uuid4()) collection_name = f"collection_{collection_id}" vectorstore = document_retrieval.create_vector_store(text_chunks, embeddings, output_db=None, collection_name=collection_name) document_retrieval.init_retriever(vectorstore) conversation_chain = document_retrieval.get_qa_retrieval_chain() return conversation_chain, vectorstore, document_retrieval, collection_name, "Complete! You can now ask questions." except Exception as e: return conversation_chain, vectorstore, document_retrieval, collection_name, f"An error occurred while processing: {str(e)}" caution_text = """⚠️ Note: depending on the size of your document, this could take several minutes. """ with gr.Blocks() as demo: vectorstore = gr.State() conversation_chain = gr.State() document_retrieval = gr.State() collection_name=gr.State() gr.Markdown("# Enterprise Knowledge Retriever", elem_id="title") gr.Markdown("Powered by LLama3.1-8B-Instruct on SambaNova Cloud. Get your API key [here](https://cloud.sambanova.ai/apis).") api_key = gr.Textbox(label="API Key", type="password", placeholder="(Optional) Enter your API key here for more availability") # Step 1: Add PDF file gr.Markdown("## 1️⃣ Upload PDF") docs = gr.File(label="Add PDF file (single)", file_types=["pdf"], file_count="single") # Step 2: Process PDF file gr.Markdown(("## 2️⃣ Process document and create vector store")) db_btn = gr.Radio(["ChromaDB"], label="Vector store type", value = "ChromaDB", type="index", info="Choose your vector store") setup_output = gr.Textbox(label="Processing status", visible=True, value="None") process_btn = gr.Button("🔄 Process") gr.Markdown(caution_text) # Preprocessing events process_btn.click(process_documents, inputs=[docs, collection_name, document_retrieval, vectorstore, conversation_chain, api_key], outputs=[conversation_chain, vectorstore, document_retrieval, collection_name, setup_output], concurrency_limit=20) # Step 3: Chat with your data gr.Markdown("## 3️⃣ Chat with your document") chatbot = gr.Chatbot(label="Chatbot", show_label=True, show_share_button=False, show_copy_button=True, likeable=True) msg = gr.Textbox(label="Ask questions about your data", show_label=True, placeholder="Enter your message...") clear_btn = gr.Button("Clear chat") sources_output = gr.Textbox(label="Sources", visible=False) # Chatbot events msg.submit(handle_userinput, inputs=[msg, conversation_chain, chatbot], outputs=[chatbot, msg], queue=False) clear_btn.click(lambda: [None, ""], inputs=None, outputs=[chatbot, msg], queue=False) if __name__ == "__main__": demo.launch()