Spaces:
Runtime error
Runtime error
| import nest_asyncio | |
| from typing import List | |
| import streamlit as st | |
| from phi.assistant import Assistant | |
| from phi.document import Document | |
| from phi.document.reader.pdf import PDFReader | |
| from phi.document.reader.website import WebsiteReader | |
| from phi.utils.log import logger | |
| from assistant import get_auto_rag_assistant # type: ignore | |
| nest_asyncio.apply() | |
| st.set_page_config( | |
| page_title="Autonomous RAG", | |
| page_icon=":orange_heart:", | |
| ) | |
| st.title("Autonomous RAG with Llama3") | |
| # st.markdown("##### :orange_heart: built using [phidata](https://github.com/phidatahq/phidata)") | |
| def restart_assistant(): | |
| logger.debug("---*--- Restarting Assistant ---*---") | |
| st.session_state["auto_rag_assistant"] = None | |
| st.session_state["auto_rag_assistant_run_id"] = None | |
| if "url_scrape_key" in st.session_state: | |
| st.session_state["url_scrape_key"] += 1 | |
| if "file_uploader_key" in st.session_state: | |
| st.session_state["file_uploader_key"] += 1 | |
| st.rerun() | |
| def main() -> None: | |
| # Get LLM model | |
| llm_model = st.sidebar.selectbox("Select LLM", options=["llama3-70b-8192", "llama3-8b-8192"]) | |
| # Set assistant_type in session state | |
| if "llm_model" not in st.session_state: | |
| st.session_state["llm_model"] = llm_model | |
| # Restart the assistant if assistant_type has changed | |
| elif st.session_state["llm_model"] != llm_model: | |
| st.session_state["llm_model"] = llm_model | |
| restart_assistant() | |
| # Get Embeddings model | |
| embeddings_model = st.sidebar.selectbox( | |
| "Select Embeddings", | |
| options=["text-embedding-3-small", "nomic-embed-text"], | |
| help="When you change the embeddings model, the documents will need to be added again.", | |
| ) | |
| # Set assistant_type in session state | |
| if "embeddings_model" not in st.session_state: | |
| st.session_state["embeddings_model"] = embeddings_model | |
| # Restart the assistant if assistant_type has changed | |
| elif st.session_state["embeddings_model"] != embeddings_model: | |
| st.session_state["embeddings_model"] = embeddings_model | |
| st.session_state["embeddings_model_updated"] = True | |
| restart_assistant() | |
| # Get the assistant | |
| auto_rag_assistant: Assistant | |
| if "auto_rag_assistant" not in st.session_state or st.session_state["auto_rag_assistant"] is None: | |
| logger.info(f"---*--- Creating {llm_model} Assistant ---*---") | |
| auto_rag_assistant = get_auto_rag_assistant(llm_model=llm_model, embeddings_model=embeddings_model) | |
| st.session_state["auto_rag_assistant"] = auto_rag_assistant | |
| else: | |
| auto_rag_assistant = st.session_state["auto_rag_assistant"] | |
| # Create assistant run (i.e. log to database) and save run_id in session state | |
| try: | |
| st.session_state["auto_rag_assistant_run_id"] = auto_rag_assistant.create_run() | |
| except Exception: | |
| st.warning("Could not create assistant, is the database running?") | |
| return | |
| # Load existing messages | |
| assistant_chat_history = auto_rag_assistant.memory.get_chat_history() | |
| if len(assistant_chat_history) > 0: | |
| logger.debug("Loading chat history") | |
| st.session_state["messages"] = assistant_chat_history | |
| else: | |
| logger.debug("No chat history found") | |
| st.session_state["messages"] = [{"role": "assistant", "content": "Upload a doc and ask me questions..."}] | |
| # Prompt for user input | |
| if prompt := st.chat_input(): | |
| st.session_state["messages"].append({"role": "user", "content": prompt}) | |
| # Display existing chat messages | |
| for message in st.session_state["messages"]: | |
| if message["role"] == "system": | |
| continue | |
| with st.chat_message(message["role"]): | |
| st.write(message["content"]) | |
| # If last message is from a user, generate a new response | |
| last_message = st.session_state["messages"][-1] | |
| if last_message.get("role") == "user": | |
| question = last_message["content"] | |
| with st.chat_message("assistant"): | |
| resp_container = st.empty() | |
| # Streaming is not supported with function calling on Groq atm | |
| response = auto_rag_assistant.run(question, stream=False) | |
| resp_container.markdown(response) # type: ignore | |
| # Once streaming is supported, the following code can be used | |
| # response = "" | |
| # for delta in auto_rag_assistant.run(question): | |
| # response += delta # type: ignore | |
| # resp_container.markdown(response) | |
| st.session_state["messages"].append({"role": "assistant", "content": response}) | |
| # Load knowledge base | |
| if auto_rag_assistant.knowledge_base: | |
| # -*- Add websites to knowledge base | |
| if "url_scrape_key" not in st.session_state: | |
| st.session_state["url_scrape_key"] = 0 | |
| input_url = st.sidebar.text_input( | |
| "Add URL to Knowledge Base", type="default", key=st.session_state["url_scrape_key"] | |
| ) | |
| add_url_button = st.sidebar.button("Add URL") | |
| if add_url_button: | |
| if input_url is not None: | |
| alert = st.sidebar.info("Processing URLs...", icon="ℹ️") | |
| if f"{input_url}_scraped" not in st.session_state: | |
| scraper = WebsiteReader(max_links=2, max_depth=1) | |
| web_documents: List[Document] = scraper.read(input_url) | |
| if web_documents: | |
| auto_rag_assistant.knowledge_base.load_documents(web_documents, upsert=True) | |
| else: | |
| st.sidebar.error("Could not read website") | |
| st.session_state[f"{input_url}_uploaded"] = True | |
| alert.empty() | |
| restart_assistant() | |
| # Add PDFs to knowledge base | |
| if "file_uploader_key" not in st.session_state: | |
| st.session_state["file_uploader_key"] = 100 | |
| uploaded_file = st.sidebar.file_uploader( | |
| "Add a PDF :page_facing_up:", type="pdf", key=st.session_state["file_uploader_key"] | |
| ) | |
| if uploaded_file is not None: | |
| alert = st.sidebar.info("Processing PDF...", icon="🧠") | |
| rag_name = uploaded_file.name.split(".")[0] | |
| if f"{rag_name}_uploaded" not in st.session_state: | |
| reader = PDFReader() | |
| rag_documents: List[Document] = reader.read(uploaded_file) | |
| if rag_documents: | |
| auto_rag_assistant.knowledge_base.load_documents(rag_documents, upsert=True) | |
| else: | |
| st.sidebar.error("Could not read PDF") | |
| st.session_state[f"{rag_name}_uploaded"] = True | |
| alert.empty() | |
| restart_assistant() | |
| if auto_rag_assistant.knowledge_base and auto_rag_assistant.knowledge_base.vector_db: | |
| if st.sidebar.button("Clear Knowledge Base"): | |
| auto_rag_assistant.knowledge_base.vector_db.clear() | |
| st.sidebar.success("Knowledge base cleared") | |
| restart_assistant() | |
| if auto_rag_assistant.storage: | |
| auto_rag_assistant_run_ids: List[str] = auto_rag_assistant.storage.get_all_run_ids() | |
| new_auto_rag_assistant_run_id = st.sidebar.selectbox("Run ID", options=auto_rag_assistant_run_ids) | |
| if st.session_state["auto_rag_assistant_run_id"] != new_auto_rag_assistant_run_id: | |
| logger.info(f"---*--- Loading {llm_model} run: {new_auto_rag_assistant_run_id} ---*---") | |
| st.session_state["auto_rag_assistant"] = get_auto_rag_assistant( | |
| llm_model=llm_model, embeddings_model=embeddings_model, run_id=new_auto_rag_assistant_run_id | |
| ) | |
| st.rerun() | |
| if st.sidebar.button("New Run"): | |
| restart_assistant() | |
| if "embeddings_model_updated" in st.session_state: | |
| st.sidebar.info("Please add documents again as the embeddings model has changed.") | |
| st.session_state["embeddings_model_updated"] = False | |
| main() | |