ollama:
  servers:
    - url: "http://127.0.0.1:11434"
      name: "Server1"
      priority: 1
      primary: true
    - url: "http://192.168.1.99:11434"
      name: "Server2"
      priority: 2
      backup: true
    - url: "http://192.168.1.88:11434"
      name: "Server3"
      priority: 2
      backup: true

  default_server: "Server1"

import re

def parse_prompt(prompt):
    # Regular expression to match "On-Server <server_name> <task>"
    match = re.match(r"On-Server (\w+) (.+)", prompt)
    if match:
        server_name = match.group(1)
        task = match.group(2)
        return server_name, task
    return None, prompt

def run_ollama(server_name, task, model_name):
    # Get server URL from llmservers.yaml
    server_url = get_server_url(server_name)
    # Run ollama command with server URL and task
    command = f"OLLAMA_HOST={server_url} ollama run {model_name} '{task}'"
    # Execute command and return response
    return execute_command(command)

def get_server_url(server_name):
    # Load llmservers.yaml and return server URL for given server name
    with open("llmservers.yaml", "r") as f:
        config = yaml.safe_load(f)
        for server in config["ollama"]["servers"]:
            if server["name"] == server_name:
                return server["url"]
    return None

def main():
    prompt = input("Enter prompt: ")
    server_name, task = parse_prompt(prompt)
    if server_name:
        model_name = "${MODEL_NAME}"
        response = run_ollama(server_name, task, model_name)
        print(response)
    else:
        # Default server
        model_name = "${MODEL_NAME}"
        response = run_ollama("Server1", prompt, model_name)
        print(response)