Update app.py

app.py

@@ -1,4 +1,3 @@
-!pip install langchain-community # Install the missing module
 import streamlit as st
 import logging
 import os
@@ -6,6 +5,8 @@ import tempfile
 import shutil
 import pdfplumber
 import ollama
+import time
+import httpx
 
 from langchain_community.document_loaders import UnstructuredPDFLoader
 from langchain_community.embeddings import OllamaEmbeddings
@@ -35,18 +36,29 @@ logging.basicConfig(
 
 logger = logging.getLogger(__name__)
 
+def ollama_list_with_retry(retries=3, delay=5):
+    """Attempt to list models from Ollama with retry logic."""
+    for attempt in range(retries):
+        try:
+            response = ollama.list()
+            logger.info("Successfully retrieved model list from Ollama")
+            return response
+        except httpx.ConnectError as e:
+            logger.error(f"Connection error: {e}. Attempt {attempt + 1} of {retries}")
+            if attempt < retries - 1:
+                time.sleep(delay)
+            else:
+                logger.error("All retry attempts failed. Cannot connect to Ollama service.")
+                raise
 
 @st.cache_resource(show_spinner=True)
-def extract_model_names(
-    models_info: Dict[str, List[Dict[str, Any]]],
-) -> Tuple[str, ...]:
+def extract_model_names(models_info: Dict[str, List[Dict[str, Any]]]) -> Tuple[str, ...]:
     """Extract model names from the provided models information."""
     logger.info("Extracting model names from models_info")
     model_names = tuple(model["name"] for model in models_info["models"])
     logger.info(f"Extracted model names: {model_names}")
     return model_names
 
-
 def create_vector_db(file_upload) -> Chroma:
     """Create a vector database from an uploaded PDF file."""
     logger.info(f"Creating vector DB from file upload: {file_upload.name}")
@@ -73,7 +85,6 @@ def create_vector_db(file_upload) -> Chroma:
     logger.info(f"Temporary directory {temp_dir} removed")
     return vector_db
 
-
 def process_question(question: str, vector_db: Chroma, selected_model: str) -> str:
     """Process a user question using the vector database and selected language model."""
     logger.info(f"Processing question: {question} using model: {selected_model}")
@@ -113,7 +124,6 @@ def process_question(question: str, vector_db: Chroma, selected_model: str) -> s
     logger.info("Question processed and response generated")
     return response
 
-
 @st.cache_data
 def extract_all_pages_as_images(file_upload) -> List[Any]:
     """Extract all pages from a PDF file as images."""
@@ -124,7 +134,6 @@ def extract_all_pages_as_images(file_upload) -> List[Any]:
     logger.info("PDF pages extracted as images")
     return pdf_pages
 
-
 def delete_vector_db(vector_db: Optional[Chroma]) -> None:
     """Delete the vector database and clear related session state."""
     logger.info("Deleting vector DB")
@@ -140,13 +149,16 @@ def delete_vector_db(vector_db: Optional[Chroma]) -> None:
         st.error("No vector database found to delete.")
         logger.warning("Attempted to delete vector DB, but none was found")
 
-
 def main() -> None:
     """Main function to run the Streamlit application."""
     st.subheader("🧠 Ollama PDF RAG playground", divider="gray", anchor=False)
 
-    models_info = ollama.list()
-    available_models = extract_model_names(models_info)
+    try:
+        models_info = ollama_list_with_retry()
+        available_models = extract_model_names(models_info)
+    except httpx.ConnectError:
+        st.error("Could not connect to the Ollama service. Please check your setup and try again.")
+        return
 
     col1, col2 = st.columns([1.5, 2])
 
@@ -154,44 +166,4 @@ def main() -> None:
         st.session_state["messages"] = []
 
     if "vector_db" not in st.session_state:
-        st.session_state["vector_db"] = None
-
-    if available_models:
-        selected_model = col2.selectbox(
-            "Pick a model available locally on your system ↓", available_models
-        )
-
-    file_upload = col1.file_uploader(
-        "Upload a PDF file ↓", type="pdf", accept_multiple_files=False
-    )
-
-    if file_upload:
-        st.session_state["file_upload"] = file_upload
-        if st.session_state["vector_db"] is None:
-            st.session_state["vector_db"] = create_vector_db(file_upload)
-        pdf_pages = extract_all_pages_as_images(file_upload)
-        st.session_state["pdf_pages"] = pdf_pages
-
-        zoom_level = col1.slider(
-            "Zoom Level", min_value=100, max_value=1000, value=700, step=50
-        )
-
-        with col1:
-            with st.container(height=410, border=True):
-                for page_image in pdf_pages:
-                    st.image(page_image, width=zoom_level)
-
-    delete_collection = col1.button("⚠️ Delete collection", type="secondary")
-
-    if delete_collection:
-        delete_vector_db(st.session_state["vector_db"])
-
-    with col2:
-        message_container = st.container(height=500, border=True)
-
-        for message in st.session_state["messages"]:
-            avatar = "🤖" if message["role"] == "assistant" else "😎"
-            with message_container.chat_message(message["role"], avatar=avatar):
-                st.markdown(message["content"])
-
-       
+        st.session_state["vector_db"]