Update app.py

app.py

@@ -42,38 +42,41 @@ def get_page_urls(url):
 
 
 
-def get_pdf_content(url, response):
-    pdf = io.BytesIO(response.content)
-    doc = fitz.open(stream=pdf, filetype="pdf")
-    return ''.join(page.get_text("text") for page in doc)
-
-def get_html_content(url, response):
-    soup = BeautifulSoup(response.content, 'html.parser')
-    content = soup.find_all('div', class_='wpb_content_element')
-    return ' '.join(c.get_text().strip() for c in content if c.get_text().strip())
-
 def get_url_content(url):
-    response = requests.get(url)
-    response.raise_for_status() 
-    if response.ok:
+    try:
+        response = requests.get(url)
+        response.raise_for_status()
         if url.endswith('.pdf'):
-            return Document(text=get_pdf_content(url, response), tags={'url': url})
+            pdf = io.BytesIO(response.content)
+            doc = fitz.open(stream=pdf, filetype="pdf")
+            text = ''.join([page.get_text("text") for page in doc])
         else:
-            return Document(text=get_html_content(url, response), tags={'url': url})
-    else:
-        st.error(f"Failed to process URL content: {e}")
-        return None
+            soup = BeautifulSoup(response.content, 'html.parser')
+            content = soup.find_all('div', class_='wpb_content_element')
+            text = ' '.join([c.get_text().strip() for c in content if c.get_text().strip() != ''])
 
+        # Create a single document with metadata
+        return DocumentArray([{'text': text, 'tags': {'url': url}}])
+    except Exception as e:
+        st.error(f"Failed to process URL content: {e}")
+        return DocumentArray()
 
 @st.cache_resource
 def get_retriever(urls):
-    documents = DocumentArray([doc for url in urls if (doc := get_url_content(url)) is not None])
-
+    documents = DocumentArray()
+    for url in urls:
+        content = get_url_content(url)
+        if content:
+            documents.extend(content)
+    
     model = SentenceTransformer('all-MiniLM-L6-v2')
-    documents.embeddings = model.encode([doc.text for doc in documents], show_progress_bar=True)
+    embeddings = model.encode([doc.text for doc in documents], show_progress_bar=True)
+    for doc, emb in zip(documents, embeddings):
+        doc.embedding = emb
     
     return documents
 
+
 @st.cache_resource
 def create_chain(_retriever):
     # A stream handler to direct streaming output on the chat screen.