Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import subprocess
-import time
 import requests
+import time
 import logging
 from langchain_community.llms import Ollama
 from langchain.callbacks.manager import CallbackManager
@@ -10,11 +10,11 @@ from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Global cache to store loaded models
+# Cache for loaded models
 loaded_models = {}
 
-# Function to check if Ollama is running
 def check_ollama_running():
+    """Wait until Ollama is fully ready."""
     url = "http://127.0.0.1:11434/api/tags"
     for _ in range(10):  # Try for ~10 seconds
         try:
@@ -23,46 +23,40 @@ def check_ollama_running():
                 logger.info("Ollama is running.")
                 return True
         except requests.exceptions.RequestException:
-            logger.warning("Ollama is not running yet. Retrying...")
-        time.sleep(1)
+            logger.warning("Waiting for Ollama to start...")
+        time.sleep(2)
     raise RuntimeError("Ollama is not running. Please check the server.")
 
-# Function to pull a model if not already available
 def pull_model(model_name):
+    """Ensure the model is available before use."""
+    if model_name in loaded_models:
+        logger.info(f"Model {model_name} is already loaded.")
+        return
     try:
-        logger.info(f"Pulling model: {model_name}")
         subprocess.run(["ollama", "pull", model_name], check=True)
         logger.info(f"Model {model_name} pulled successfully.")
+        loaded_models[model_name] = True
     except subprocess.CalledProcessError as e:
         logger.error(f"Failed to pull model {model_name}: {e}")
         raise
 
-# Function to get an LLM instance with streaming enabled
 def get_llm(model_name):
+    """Get an LLM instance with streaming enabled."""
     callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
     return Ollama(model=model_name, base_url="http://127.0.0.1:11434", callback_manager=callback_manager)
 
-# Function to check and load a model
-def check_and_load_model(model_name):
-    if model_name in loaded_models:
-        logger.info(f"Model {model_name} is already loaded.")
-        return loaded_models[model_name]
-    pull_model(model_name)  # Ensure the model is available
-    llm = get_llm(model_name)
-    loaded_models[model_name] = llm
-    return llm
-
-# Function to handle Gradio input with streaming
 def query_model(model_name, prompt):
-    check_ollama_running()  # Ensure Ollama is running before making requests
-    llm = check_and_load_model(model_name)
+    """Generate responses from the model with streaming."""
+    check_ollama_running()  # Ensure Ollama is ready
+    pull_model(model_name)  # Make sure the model is available
+    llm = get_llm(model_name)  # Load the model
 
     response = ""
     for token in llm.stream(prompt):
         response += token
-        yield response  # Stream the response to Gradio in real-time
+        yield response  # Stream response in real-time
 
-# Define the Gradio interface
+# Define Gradio interface
 iface = gr.Interface(
     fn=query_model,
     inputs=[
@@ -76,4 +70,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=8080)
+    iface.launch(server_name="0.0.0.0", server_port=7860)