diff --git a/start_server.py b/start_server.py index 1dba52e..1a89e16 100644 --- a/start_server.py +++ b/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()