Update app.py

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 from bs4 import BeautifulSoup
 import io
-import fitz
+import fitz # PyMuPDF
 import requests
 from langchain.llms import LlamaCpp
 from langchain.callbacks.base import BaseCallbackHandler
@@ -42,20 +42,25 @@ 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)
-    if url.endswith('.pdf'):
-        pdf = io.BytesIO(response.content)
-        doc = fitz.open(stream=pdf, filetype="pdf")
-        text = ''.join([page.get_text("text") for page in doc])
-        return Document(text=text, tags={'url': url})
+    if response.ok:
+        if url.endswith('.pdf'):
+            return Document(text=get_pdf_content(url, response), tags={'url': url})
+        else:
+            return Document(text=get_html_content(url, response), tags={'url': url})
     else:
-        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() != ''])
-        return Document(text=text, tags={'url': url})
-
-
+        raise ValueError("Failed to fetch URL content")