Update app.py

app.py

@@ -1,21 +1,77 @@
 import gradio as gr
-from langchain_ollama import OllamaLLM
+import subprocess
+import time
+import requests
+import logging
+from langchain_community.llms import Ollama
+from langchain.callbacks.manager import CallbackManager
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 
-llm = OllamaLLM(model="deepseek-r1:1.5b")
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-def query_model(prompt: str) -> str:
-    """
-    Uses LangChain's Ollama integration to process the input prompt.
-    """
-    response = llm(prompt)
-    return response
+# Global cache to store loaded models
+loaded_models = {}
 
+# Function to check if Ollama is running
+def check_ollama_running():
+    url = "http://127.0.0.1:11434/api/tags"
+    for _ in range(10):  # Try for ~10 seconds
+        try:
+            response = requests.get(url, timeout=2)
+            if response.status_code == 200:
+                logger.info("Ollama is running.")
+                return True
+        except requests.exceptions.RequestException:
+            logger.warning("Ollama is not running yet. Retrying...")
+        time.sleep(1)
+    raise RuntimeError("Ollama is not running. Please check the server.")
+
+# Function to pull a model if not already available
+def pull_model(model_name):
+    try:
+        logger.info(f"Pulling model: {model_name}")
+        subprocess.run(["ollama", "pull", model_name], check=True)
+        logger.info(f"Model {model_name} pulled successfully.")
+    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):
+    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)
+
+    response = ""
+    for token in llm.stream(prompt):
+        response += token
+        yield response  # Stream the response to Gradio in real-time
+
+# Define the Gradio interface
 iface = gr.Interface(
     fn=query_model,
-    inputs="text",
+    inputs=[
+        gr.Dropdown(["deepseek-r1:1.5b", "mistral:7b"], label="Select Model"),
+        gr.Textbox(label="Enter your prompt")
+    ],
     outputs="text",
     title="Ollama via LangChain & Gradio",
-    description="Enter a prompt to interact with the Ollama-based model.",
+    description="Enter a prompt to interact with the Ollama-based model with streaming response.",
     flagging_dir="/app/flagged"
 )