Add ollama check to start server

start_server.py

@@ -1,10 +1,72 @@
 import uvicorn
 import argparse
+import sys
+import tomli
+import requests
+import logging
 from dotenv import load_dotenv
 from core.config import get_settings
 from core.logging_config import setup_logging
 
 
+def check_ollama_running(base_url):
+    """Check if Ollama is running and accessible at the given URL."""
+    try:
+        api_url = f"{base_url}/api/tags"
+        response = requests.get(api_url, timeout=2)
+        return response.status_code == 200
+    except requests.RequestException:
+        return False
+
+
+def get_ollama_usage_info():
+    """Check if Ollama is required based on the configuration file and get base URLs."""
+    try:
+        with open("databridge.toml", "rb") as f:
+            config = tomli.load(f)
+
+        ollama_configs = []
+
+        # Check embedding provider
+        if "embedding" in config and config["embedding"].get("provider") == "ollama":
+            base_url = config["embedding"].get("base_url")
+            if base_url:
+                ollama_configs.append({"component": "embedding", "base_url": base_url})
+
+        # Check completion provider
+        if "completion" in config and config["completion"].get("provider") == "ollama":
+            base_url = config["completion"].get("base_url")
+            if base_url:
+                ollama_configs.append({"component": "completion", "base_url": base_url})
+
+        # Check rules provider
+        if "rules" in config and config["rules"].get("provider") == "ollama":
+            base_url = config["rules"].get("base_url")
+            if base_url:
+                ollama_configs.append({"component": "rules", "base_url": base_url})
+
+        # Check graph provider
+        if "graph" in config and config["graph"].get("provider") == "ollama":
+            base_url = config["graph"].get("base_url")
+            if base_url:
+                ollama_configs.append({"component": "graph", "base_url": base_url})
+
+        # Check parser.vision provider
+        if (
+            "parser" in config
+            and "vision" in config["parser"]
+            and config["parser"]["vision"].get("provider") == "ollama"
+        ):
+            base_url = config["parser"]["vision"].get("base_url")
+            if base_url:
+                ollama_configs.append({"component": "parser.vision", "base_url": base_url})
+
+        return ollama_configs
+    except Exception as e:
+        logging.error(f"Error checking Ollama configuration: {e}")
+        return []
+
+
 def main():
     # Parse command line arguments
     parser = argparse.ArgumentParser(description="Start the DataBridge server")
@@ -14,6 +76,11 @@ def main():
         default="info",
         help="Set the logging level",
     )
+    parser.add_argument(
+        "--skip-ollama-check",
+        action="store_true",
+        help="Skip Ollama availability check",
+    )
     args = parser.parse_args()
 
     # Set up logging first with specified level
@@ -22,6 +89,35 @@ def main():
     # Load environment variables from .env file
     load_dotenv()
 
+    # Check if Ollama is required and running
+    if not args.skip_ollama_check:
+        ollama_configs = get_ollama_usage_info()
+
+        if ollama_configs:
+            # Group configs by base_url to avoid redundant checks
+            base_urls = {}
+            for config in ollama_configs:
+                if config["base_url"] not in base_urls:
+                    base_urls[config["base_url"]] = []
+                base_urls[config["base_url"]].append(config["component"])
+
+            all_running = True
+            for base_url, components in base_urls.items():
+                if not check_ollama_running(base_url):
+                    print(f"ERROR: Ollama is not accessible at {base_url}")
+                    print(f"This URL is used by these components: {', '.join(components)}")
+                    all_running = False
+
+            if not all_running:
+                print(
+                    "\nPlease ensure Ollama is running at the configured URLs before starting the server"
+                )
+                print("Run with --skip-ollama-check to bypass this check")
+                sys.exit(1)
+            else:
+                component_list = [config["component"] for config in ollama_configs]
+                print(f"Ollama is running and will be used for: {', '.join(component_list)}")
+
     # Load settings (this will validate all required env vars)
     settings = get_settings()