morphik-core/start_server.py

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")
    parser.add_argument(
        "--log",
        choices=["debug", "info", "warning", "error"],
        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
    setup_logging(log_level=args.log.upper())

    # 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()

    # Start server
    uvicorn.run(
        "core.api:app",
        host=settings.HOST,
        port=settings.PORT,
        loop="asyncio",
        log_level=args.log,
        # reload=settings.RELOAD
    )


if __name__ == "__main__":
    main()