import streamlit as st from jamaibase import JamAI, protocol as p import time import os # Initialize session state variables def initialize_session_state(): if "unique_time" not in st.session_state: st.session_state.unique_time = time.time() if "knowledge_base_exist" not in st.session_state: st.session_state.knowledge_base_exist = False if "chat_history" not in st.session_state: st.session_state.chat_history = [] if "api_key" not in st.session_state: st.session_state.api_key = None if "project_id" not in st.session_state: st.session_state.project_id = None if "model" not in st.session_state: st.session_state.model = "ellm/Qwen/Qwen2-7B-Instruct" if "k" not in st.session_state: st.session_state.k = 2 if "temperature" not in st.session_state: st.session_state.temperature = 0.01 if "top_p" not in st.session_state: st.session_state.top_p = 0.01 if "max_tokens" not in st.session_state: st.session_state.max_tokens = 496 # Function to check if any input is None def check_all_fields_filled(file_upload, api_key, proj_id): return file_upload is not None and api_key and proj_id def clear_credentials(): st.session_state.api_key = "" st.session_state.project_id = "" st.session_state.unique_time = time.time() # reset the unique time # Function to create knowledge base def create_knowledge_base(jamai, file_upload): try: with st.spinner("Creating Knowledge Base..."): knowledge_simple = f"knowledge-simple-{st.session_state.unique_time}" knowledge_table = jamai.create_knowledge_table( p.KnowledgeTableSchemaCreate( id=knowledge_simple, cols=[], embedding_model="ellm/BAAI/bge-m3", ) ) st.success("Successfully created Knowledge Base") # Save PDF to local directory current_dir = os.path.dirname(os.path.abspath(__file__)) file_path = os.path.join(current_dir, file_upload.name) with open(file_path, "wb") as f: f.write(file_upload.read()) # Upload file to knowledge base with st.spinner("Uploading PDF to Knowledge Base..."): response = jamai.upload_file( p.FileUploadRequest( file_path=file_path, table_id=knowledge_simple, ) ) assert response.ok st.success("Successfully uploaded PDF to Knowledge Base!") st.session_state.knowledge_base_exist = True # remove after upload successfully os.remove(file_path) return knowledge_simple except Exception as e: clear_credentials() # st.error(f"An error occurred: {str(e)}") st.warning("An error occurred. Please check your credentials and try again.") return None # Function to create chat table def create_chat_table(jamai, knowledge_simple): try: with st.spinner("Creating Chat Table..."): table = jamai.create_chat_table( p.ChatTableSchemaCreate( id=f"chat-rag-{st.session_state.unique_time}", cols=[ p.ColumnSchemaCreate(id="User", dtype=p.DtypeCreateEnum.str_), p.ColumnSchemaCreate( id="AI", dtype=p.DtypeCreateEnum.str_, gen_config=p.ChatRequest( model=st.session_state.model, messages=[p.ChatEntry.system("You are a concise assistant.")], rag_params=p.RAGParams( table_id=knowledge_simple, k=st.session_state.k, ), temperature=st.session_state.temperature, top_p=st.session_state.top_p, max_tokens=st.session_state.max_tokens, ).model_dump(), ), ], ) ) st.success("Successfully created Chat Table") except Exception as e: clear_credentials() # st.error(f"An error occurred while creating the chat table: {str(e)}") st.warning("An error occurred. Please check your credentials and try again.") # Function to ask a question with improved streaming output def ask_question(question): jamai = JamAI(api_key=st.session_state.api_key, project_id=st.session_state.project_id) completion = jamai.add_table_rows( "chat", p.RowAddRequest( table_id=f"chat-rag-{st.session_state.unique_time}", data=[dict(User=question)], stream=True, ), ) full_response = "" for chunk in completion: if chunk.output_column_name != "AI": continue if isinstance(chunk, p.GenTableStreamReferences): pass else: full_response += chunk.text yield full_response # Main app function def main(): st.title("Chat With PDF") initialize_session_state() # Sidebar for configuration with st.sidebar: st.logo("resource/Jamai-Long-Black-Main.icuEAbYB.svg") login_button = """ Login to your Jamai Base """ st.markdown(login_button, unsafe_allow_html=True) st.header("Configuration") file_upload = st.file_uploader("Upload your PDF", type=["pdf"], disabled=st.session_state.knowledge_base_exist) api_key = st.text_input('JAMAI API KEY', value=st.session_state.get('api_key', ''), type='password', disabled=st.session_state.knowledge_base_exist) project_id = st.text_input('Project ID', value=st.session_state.get('project_id', ''), disabled=st.session_state.knowledge_base_exist) with st.expander("Advanced Settings"): model_options = [ 'ellm/Qwen/Qwen2-7B-Instruct', 'ellm/Qwen/Qwen2-72B-Instruct', 'ellm/meta-llama/Llama-3-8B-Instruct', 'ellm/meta-llama/Llama-3-70B-Instruct', 'ellm/meta-llama/Llama-3.1-8B-Instruct', 'ellm/meta-llama/Llama-3.1-70B-Instruct', 'ellm/microsoft/Phi-3-mini-128k-Instruct-Int4' 'ellm/microsoft/Phi-3-medium-128k-Instruct-Int4', ] st.session_state.model = st.selectbox("Model", options=model_options, index=0, disabled=st.session_state.knowledge_base_exist) st.session_state.k = st.slider("k", value=2, min_value=1, max_value=30, step=1, disabled=st.session_state.knowledge_base_exist) st.session_state.max_tokens = st.slider("max tokens", value=496, min_value=96, max_value=960, step=8, disabled=st.session_state.knowledge_base_exist) temperature_options = [str(i/10) for i in range(1,11)] st.session_state.temperature = st.selectbox("temperature", options=temperature_options, format_func=lambda x: float(x), disabled=st.session_state.knowledge_base_exist) top_p_options = [str(i/10) for i in range(1,11)] st.session_state.top_p = st.selectbox("top p", options=top_p_options, format_func=lambda x: float(x), disabled=st.session_state.knowledge_base_exist) if st.button("Create Knowledge Base", disabled=st.session_state.knowledge_base_exist): if not check_all_fields_filled(file_upload, api_key, project_id): st.warning("Please provide all required information: PDF file, JAMAI API KEY, and Project ID.") else: st.session_state.api_key = api_key st.session_state.project_id = project_id jamai = JamAI(api_key=api_key, project_id=project_id) knowledge_simple = create_knowledge_base(jamai, file_upload) if knowledge_simple: create_chat_table(jamai, knowledge_simple) # Main chat interface st.header("Jamai Base") # Display chat history for message in st.session_state.chat_history: with st.chat_message(message["role"]): st.markdown(message["content"]) # Chat input if question := st.chat_input("Ask a question about your PDF"): if st.session_state.knowledge_base_exist: st.session_state.chat_history.append({"role": "user", "content": question}) with st.chat_message("user"): st.markdown(question) with st.chat_message("assistant"): message_placeholder = st.empty() full_response = "" for response in ask_question(question): message_placeholder.markdown(response + "▌") full_response = response message_placeholder.markdown(full_response) st.session_state.chat_history.append({"role": "assistant", "content": full_response}) else: st.warning("Please upload a PDF and create a Knowledge Base first.") if __name__ == "__main__": main()