Update app.py

app.py

@@ -1,7 +1,191 @@
-from fastapi import FastAPI
+from typing import List
 
-app = FastAPI()
+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
 
-@app.get("/")
-def greet_json():
-    return {"Hello": "World!"}
+from assistant import get_groq_assistant  # type: ignore
+
+st.set_page_config(
+    page_title="ISW RAG",
+    page_icon=":books:",
+)
+st.title("RAG with Llama3 on Groq")
+st.markdown("Built at ISW")
+
+import os
+
+from groq import Groq
+
+client = Groq(
+    api_key=os.environ.get("GROQ_API_KEY"),
+)
+
+chat_completion = client.chat.completions.create(
+    messages=[
+        {
+            "role": "user",
+            "content": "Explain the importance of fast language models",
+        }
+    ],
+    model="llama3-8b-8192",
+)
+
+print(chat_completion.choices[0].message.content)
+
+print(chat_completion.choices[0].message.content)
+
+def restart_assistant():
+    st.session_state["rag_assistant"] = None
+    st.session_state["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", "mixtral-8x7b-32768"])
+    # 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=["nomic-embed-text", "text-embedding-3-small"],
+        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
+    rag_assistant: Assistant
+    if "rag_assistant" not in st.session_state or st.session_state["rag_assistant"] is None:
+        logger.info(f"---*--- Creating {llm_model} Assistant ---*---")
+        rag_assistant = get_groq_assistant(llm_model=llm_model, embeddings_model=embeddings_model)
+        st.session_state["rag_assistant"] = rag_assistant
+    else:
+        rag_assistant = st.session_state["rag_assistant"]
+
+    # Create assistant run (i.e. log to database) and save run_id in session state
+    try:
+        st.session_state["rag_assistant_run_id"] = rag_assistant.create_run()
+    except Exception:
+        st.warning("Could not create assistant, is the database running?")
+        return
+
+    # Load existing messages
+    assistant_chat_history = 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"):
+            response = ""
+            resp_container = st.empty()
+            for delta in 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 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:
+                        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()
+
+        # 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:
+                    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()
+
+    if rag_assistant.knowledge_base and rag_assistant.knowledge_base.vector_db:
+        if st.sidebar.button("Clear Knowledge Base"):
+            rag_assistant.knowledge_base.vector_db.clear()
+            st.sidebar.success("Knowledge base cleared")
+
+    if rag_assistant.storage:
+        rag_assistant_run_ids: List[str] = rag_assistant.storage.get_all_run_ids()
+        new_rag_assistant_run_id = st.sidebar.selectbox("Run ID", options=rag_assistant_run_ids)
+        if st.session_state["rag_assistant_run_id"] != new_rag_assistant_run_id:
+            logger.info(f"---*--- Loading {llm_model} run: {new_rag_assistant_run_id} ---*---")
+            st.session_state["rag_assistant"] = get_groq_assistant(
+                llm_model=llm_model, embeddings_model=embeddings_model, run_id=new_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()