Update app.py

app.py

@@ -1,21 +1,19 @@
 import streamlit as st
 from bs4 import BeautifulSoup
 import io
-import fitz
+import fitz  # PyMuPDF
 import requests
+from docarray import Document
+from pydantic import BaseModel, Field
+from typing import List
 from langchain.llms import LlamaCpp
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.vectorstores import DocArrayInMemorySearch
-from langchain.docstore.document import Document
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationalRetrievalChain
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
-
-# StreamHandler to intercept streaming output from the LLM.
-# This makes it appear that the Language Model is "typing"
-# in realtime.
 class StreamHandler(BaseCallbackHandler):
     def __init__(self, container, initial_text=""):
         self.container = container
@@ -25,24 +23,25 @@ class StreamHandler(BaseCallbackHandler):
         self.text += token
         self.container.markdown(self.text)
 
+class DocArrayDoc(BaseModel):
+    text: str = Field(default="")
+    embedding: List[float]
+    metadata: dict = Field(default_factory=dict)
 
 @st.cache_data
-
 def get_page_urls(url):
     page = requests.get(url)
     soup = BeautifulSoup(page.content, 'html.parser')
-    links = [link['href'] for link in soup.find_all('a') if 'href' in link.attrs and link['href'].startswith(url) and link['href'] not in [url]]
+    links = [link['href'] for link in soup.find_all('a', href=True) if link['href'].startswith(url)]
     links.append(url)
     return set(links)
 
 @st.cache(allow_output_mutation=True)
 def process_pdf(file):
-    # file is expected to be a BytesIO object directly from the file uploader
-    doc = fitz.open("pdf", file.read())  # "pdf" indicates file format is PDF, reading the BytesIO stream
+    doc = fitz.open("pdf", file.read())
     texts = [page.get_text() for page in doc]
     return '\n'.join(texts)
 
-
 def get_url_content(url):
     response = requests.get(url)
     if url.endswith('.pdf'):
@@ -52,93 +51,49 @@ def get_url_content(url):
     else:
         soup = BeautifulSoup(response.content, 'html.parser')
         content = soup.find_all('div', class_='wpb_content_element')
-        text = [c.get_text().strip() for c in content if c.get_text().strip() != '']
-        text = [line for item in text for line in item.split('\n') if line.strip() != '']
-        # Exclude footer content
-        try:
-            arts_on_index = text.index('ARTS ON:')
-            return (url, '\n'.join(text[:arts_on_index]))
-        except ValueError:
-            return (url, '\n'.join(text))  # Return full text if specific marker not found
+        text = ' '.join([c.get_text().strip() for c in content])
+        return (url, text)
 
 @st.cache_resource
 def get_retriever(urls):
     all_content = [get_url_content(url) for url in urls]
-    print(all_content)  # See what is actually fetched
-    documents = [Document(page_content=doc, metadata={'url': url}) for (url, doc) in all_content]
-    print(documents)  # Verify that documents are created correctly
-
+    documents = [Document(text=content, metadata={'url': url}) for (url, content) in all_content]
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=200)
     docs = text_splitter.split_documents(documents)
-    print(docs)  # Check the final structure of split documents
-
     embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
     db = DocArrayInMemorySearch.from_documents(docs, embeddings)
     retriever = db.as_retriever(search_type="mmr", search_kwargs={"k": 5, "fetch_k": 10})
     return retriever
 
-
 @st.cache_resource
 def create_chain(_retriever):
-    # A stream handler to direct streaming output on the chat screen.
-    # This will need to be handled somewhat differently.
-    # But it demonstrates what potential it carries.
-    # stream_handler = StreamHandler(st.empty())
-
-    # Callback manager is a way to intercept streaming output from the
-    # LLM and take some action on it. Here we are giving it our custom
-    # stream handler to make it appear as if the LLM is typing the
-    # responses in real time.
-    # callback_manager = CallbackManager([stream_handler])
-
-    n_gpu_layers = 40  # Change this value based on your model and your GPU VRAM pool.
-    n_batch = 2048  # Should be between 1 and n_ctx, consider the amount of VRAM in your GPU.
-
+    n_gpu_layers = 10
+    n_batch = 2048
     llm = LlamaCpp(
-            model_path="models /mistral-7b-instruct-v0.1.Q5_0.gguf",
+            model_path="models/mistral-7b-instruct-v0.1.Q5_0.gguf",
             n_gpu_layers=n_gpu_layers,
             n_batch=n_batch,
             n_ctx=2048,
-            # max_tokens=2048,
             temperature=0,
-            # callback_manager=callback_manager,
             verbose=False,
             streaming=True,
             )
-
-    # Template for the prompt.
-    # template = "{question}"
-
-    # We create a prompt from the template so we can use it with langchain
-    # prompt = PromptTemplate(template=template, input_variables=["question"])
-
-    # Setup memory for contextual conversation
     memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-
-    # We create a qa chain with our llm, retriever, and memory
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm, retriever=_retriever, memory=memory, verbose=False
     )
-
     return qa_chain
 
-
-
-# Set the webpage title
+# Webpage title and header
 st.set_page_config(page_title="Your own AI-Chat!")
 st.header("Your own AI-Chat!")
 
-# This sets the LLM's personality.
-# The initial personality privided is basic.
-# Try something interesting and notice how the LLM responses are affected.
-# system_prompt = st.text_area(
-#    label="System Prompt",
-#    value="You are a helpful AI assistant who answers questions in short sentences.",
-#    key="system_prompt")
-# Choose input method
+system_prompt = st.text_area(
+    label="System Prompt",
+    value="You are a helpful AI assistant who answers questions accurately.",
+    key="system_prompt")
 
 input_type = st.radio("Choose an input method:", ['URL', 'Upload PDF'])
-
 if input_type == 'URL':
     base_url = st.text_input("Enter the site URL here:", key="base_url")
     if base_url:
@@ -149,36 +104,21 @@ elif input_type == 'Upload PDF':
     uploaded_file = st.file_uploader("Upload your PDF here:", type="pdf")
     if uploaded_file:
         pdf_text = process_pdf(uploaded_file)
-        # Process the PDF text into a format that can be used by your LLM
-        urls = [pdf_text]  # Adapt as needed for your system
-        retriever = get_retriever(urls)  # Ensure your retriever can handle raw text; if not, adapt it.
+        urls = [pdf_text]  # Assuming this needs to be wrapped into proper structure
+        retriever = get_retriever(urls)  # Ensure retriever accepts this
         llm_chain = create_chain(retriever)
 
-    # We store the conversation in the session state.
-    # This will be used to render the chat conversation.
-    # We initialize it with the first message we want to be greeted with
-    
-# Initialize chat session state for storing messages and responses
-if "messages" not in st.session_state:
-    st.session_state.messages = [{"role": "assistant", "content": "How may I help you today?"}]
-
-if "current_response" not in st.session_state:
-    st.session_state.current_response = ""
-
-# Render the chat messages
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-# Input and response handling
-if llm_chain and (user_prompt := st.chat_input("Your message here", key="user_input")):
-    # Add user input to the session state and chat window
-    st.session_state.messages.append({"role": "user", "content": user_prompt})
-    with st.chat_message("user"):
-        st.markdown(user_prompt)
-    
-    # Generate and display the response using the LLM chain
-    response = llm_chain.run(user_prompt)
-    st.session_state.messages.append({"role": "assistant", "content": response})
-    with st.chat_message("assistant"):
-        st.markdown(response)
+# Interaction and message handling
+if 'retriever' in locals() and retriever:
+    if "messages" not in st.session_state:
+        st.session_state.messages = [{"role": "assistant", "content": "How may I help you today?"}]
+    if "current_response" not in st.session_state:
+        st.session_state.current_response = ""
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+    user_prompt = st.chat_input("Your message here", key="user_input")
+    if user_prompt:
+        st.session_state.messages.append({"role": "user", "content": user_prompt})
+        response = llm_chain.run(user_prompt)
+        st.session_state.messages.append({"role": "assistant", "content": response})