homelab-codex-ws/deploy_agent.py

from pathlib import Path
import subprocess

from ollama_client import ask
import yaml


def _build_prompt(service: str, extra: str = "") -> str:
    prompt = (
        "Output ONLY valid docker-compose YAML. No explanations. No markdown.\n\n"
        f"Generate a docker-compose file for this service: {service}"
    )
    if extra:
        prompt = f"{prompt}\n{extra}"
    return prompt


def _generate_compose(service: str, extra: str = "") -> str:
    base_prompt = (
        _build_prompt(service, extra)
    )
    retry_prompt = (
        "Previous output was invalid. Fix YAML and return valid docker-compose only."
    )

    prompt = base_prompt
    for _ in range(3):
        response = ask(prompt)
        if response.startswith("ERROR:"):
            return response
        try:
            parsed = yaml.safe_load(response)
        except yaml.YAMLError:
            prompt = f"{base_prompt}\n\n{retry_prompt}"
            continue
        if isinstance(parsed, dict) and "services" in parsed:
            return response
        prompt = f"{base_prompt}\n\n{retry_prompt}"
    return "ERROR: invalid docker-compose"


def generate_compose(service: str) -> str:
    return _generate_compose(service)


def _is_valid_compose(compose: str) -> bool:
    try:
        parsed = yaml.safe_load(compose)
    except yaml.YAMLError:
        return False
    return isinstance(parsed, dict) and "services" in parsed

def _get_compose_logs(path: Path) -> str:
    try:
        ids_result = subprocess.run(
            ["docker", "compose", "ps", "-q"],
            cwd=path,
            check=False,
            capture_output=True,
            text=True,
        )
    except Exception as exc:
        return str(exc)

    container_ids = [line.strip() for line in ids_result.stdout.splitlines() if line.strip()]
    if not container_ids:
        return ""

    parts = []
    total = 0
    limit = 2000
    for container_id in container_ids:
        try:
            logs_result = subprocess.run(
                ["docker", "logs", container_id, "--tail=50"],
                check=False,
                capture_output=True,
                text=True,
            )
        except Exception as exc:
            chunk = f"{container_id}:\n{exc}\n"
        else:
            chunk_body = logs_result.stdout or logs_result.stderr
            chunk = f"{container_id}:\n{chunk_body}\n"
        remaining = limit - total
        if remaining <= 0:
            break
        if len(chunk) > remaining:
            chunk = chunk[:remaining]
        parts.append(chunk)
        total += len(chunk)
        if total >= limit:
            break
    return "".join(parts).strip()


def _fix_compose(service: str, error: str, status: str, logs: str) -> str:
    prompt = (
        "Deployment failed.\n\n"
        f"Service: {service}\n\n"
        f"Error: {error}\n\n"
        f"Status: {status}\n\n"
        f"Logs: {logs}\n\n"
        "Fix the docker-compose YAML. Return ONLY corrected YAML."
    )
    response = ask(prompt)
    if response.startswith("ERROR:"):
        return response
    if not _is_valid_compose(response):
        return "ERROR: invalid docker-compose"
    return response


def get_service_status(path: Path) -> str:
    try:
        result = subprocess.run(
            ["docker", "compose", "ps"],
            cwd=path,
            check=False,
            capture_output=True,
            text=True,
        )
    except Exception as exc:
        return f"ERROR: {exc}"
    output = result.stdout.strip()
    if result.returncode != 0:
        error = result.stderr.strip() or output or "docker compose ps failed"
        return f"ERROR: {error}"
    return output


def _run_compose_up(path: Path) -> tuple[bool, str]:
    try:
        result = subprocess.run(
            ["docker", "compose", "up", "-d"],
            cwd=path,
            check=False,
            capture_output=True,
            text=True,
        )
    except Exception as exc:
        return False, str(exc)
    if result.returncode != 0:
        error = result.stderr.strip() or result.stdout.strip() or "docker compose up failed"
        return False, error
    return True, result.stdout.strip()


def deploy_service(service: str) -> str:
    prompt_extra = ""
    try:
        docker_ps = subprocess.run(
            ["docker", "ps"],
            check=False,
            capture_output=True,
            text=True,
        )
        docker_ps_output = docker_ps.stdout
        if ":80->" in docker_ps_output or "0.0.0.0:80-" in docker_ps_output or "[::]:80-" in docker_ps_output:
            prompt_extra = "Use a different port than 80"
    except Exception:
        docker_ps_output = ""

    compose = _generate_compose(service, prompt_extra)
    if compose.startswith("ERROR:"):
        return compose

    deployments_dir = Path("./deployments")
    target_dir = deployments_dir / service
    suffix = 1
    while target_dir.exists():
        target_dir = deployments_dir / f"{service}-{suffix}"
        suffix += 1

    target_dir.mkdir(parents=True, exist_ok=False)
    compose_path = target_dir / "docker-compose.yml"
    compose_path.write_text(compose, encoding="utf-8")

    ok, error = _run_compose_up(target_dir)
    if not ok:
        status = get_service_status(target_dir)
        logs = _get_compose_logs(target_dir)
        fixed_compose = _fix_compose(service, error, status, logs)
        if fixed_compose.startswith("ERROR:"):
            if status and not status.startswith("ERROR:"):
                return f"ERROR: {error}\n{status}"
            return f"ERROR: {error}"
        compose_path.write_text(fixed_compose, encoding="utf-8")
        ok, error = _run_compose_up(target_dir)
        if not ok:
            status = get_service_status(target_dir)
            logs = _get_compose_logs(target_dir)
            debug = f"""
ESCALATION REQUIRED

SERVICE: {service}

ERROR:
{error}

STATUS:
{status}

LOGS:
{logs}
"""
            return f"ESCALATE_TO_CODEX\n{debug}"

    status = get_service_status(target_dir)
    if status.startswith("ERROR:"):
        return status
    if "Up" not in status or "unhealthy" in status.lower():
        return f"ERROR: no running services\n{status}"

    return f"DEPLOYED: {target_dir.name}\n{status}"


if __name__ == "__main__":
    print(deploy_service("nginx"))