hobbs-iac/compose/hf-pdocker-01/voicecoach-v2/agent/tools.py

"""Homelab tools wired into the agent via @function_tool."""
import asyncio
import logging
import os
from typing import Annotated

import httpx
from livekit.agents import llm

logger = logging.getLogger(__name__)

# Homelab node config
NODES = {
    "hf-pdocker-01": "192.168.86.192",
    "hf-pdocker-02": "192.168.86.100",
    "bart": "192.168.86.167",
    "gpu": "192.168.86.27",
}
SSH_USER = "glyptic"
SSH_PASS = os.getenv("SSH_PASS", "&Rip8a115")
FARM_MANAGER_URL = "http://192.168.86.192:8888"
HA_URL = os.getenv("HA_URL", "http://192.168.86.192:8123")
HA_TOKEN = os.getenv("HA_TOKEN", "")
ALERTMANAGER_URL = "http://192.168.86.192:9093"
SEARXNG_URL = "http://192.168.86.100:8888"


async def _ssh_exec(host: str, cmd: str, timeout: float = 15.0) -> str:
    """Run a command on a remote host via sshpass."""
    proc = await asyncio.create_subprocess_exec(
        "sshpass", "-p", SSH_PASS,
        "ssh", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=5",
        f"{SSH_USER}@{host}", cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    try:
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
    except asyncio.TimeoutError:
        proc.kill()
        return f"ERROR: command timed out after {timeout}s"
    if proc.returncode != 0:
        return f"ERROR: {stderr.decode().strip()}"
    return stdout.decode().strip()


@llm.function_tool
async def infra_status(
    node: Annotated[str, "Node name: hf-pdocker-01, hf-pdocker-02, bart, gpu, or 'all'"],
) -> str:
    """Check container health on a homelab node. Returns running container count and any unhealthy containers."""
    nodes = list(NODES.keys()) if node == "all" else [node]
    results = []
    for n in nodes:
        ip = NODES.get(n)
        if not ip:
            results.append(f"{n}: unknown node")
            continue
        out = await _ssh_exec(ip, "docker ps --format '{{.Names}}\\t{{.Status}}' | head -50")
        lines = [l for l in out.splitlines() if l.strip()]
        unhealthy = [l for l in lines if "unhealthy" in l.lower() or "exited" in l.lower()]
        results.append(
            f"{n}: {len(lines)} containers running"
            + (f", ISSUES: {'; '.join(l.split(chr(9))[0] for l in unhealthy)}" if unhealthy else ", all healthy")
        )
    return "\n".join(results)


@llm.function_tool
async def container_action(
    node: Annotated[str, "Node name: hf-pdocker-01, hf-pdocker-02, bart, gpu"],
    container: Annotated[str, "Container name"],
    action: Annotated[str, "Action: restart, stop, or start"],
) -> str:
    """Restart, stop, or start a Docker container on a homelab node."""
    ip = NODES.get(node)
    if not ip:
        return f"Unknown node: {node}"
    if action not in ("restart", "stop", "start"):
        return f"Invalid action: {action}. Use restart, stop, or start."
    result = await _ssh_exec(ip, f"docker {action} {container}")
    return f"{action} {container} on {node}: {'done' if container in result or result == '' else result}"


@llm.function_tool
async def check_logs(
    node: Annotated[str, "Node name"],
    container: Annotated[str, "Container name"],
    lines: Annotated[int, "Number of log lines to return (max 50)"] = 20,
) -> str:
    """Get recent logs from a container."""
    ip = NODES.get(node)
    if not ip:
        return f"Unknown node: {node}"
    lines = min(lines, 50)
    return await _ssh_exec(ip, f"docker logs --tail {lines} {container} 2>&1")


@llm.function_tool
async def read_alerts() -> str:
    """Check Prometheus Alertmanager for active alerts."""
    try:
        async with httpx.AsyncClient(timeout=5.0) as client:
            r = await client.get(f"{ALERTMANAGER_URL}/api/v2/alerts?active=true&silenced=false")
            r.raise_for_status()
            alerts = r.json()
        if not alerts:
            return "No active alerts."
        parts = []
        for a in alerts[:10]:
            name = a.get("labels", {}).get("alertname", "unknown")
            severity = a.get("labels", {}).get("severity", "")
            summary = a.get("annotations", {}).get("summary", "")
            parts.append(f"[{severity}] {name}: {summary}" if summary else f"[{severity}] {name}")
        return "\n".join(parts)
    except Exception as e:
        return f"Could not reach Alertmanager: {e}"


@llm.function_tool
async def home_assistant(
    action: Annotated[str, "get_states, turn_on, turn_off, or toggle"],
    entity_id: Annotated[str, "HA entity ID, e.g. light.living_room or 'all' for get_states"] = "all",
) -> str:
    """Interact with Home Assistant — check sensor states or control devices."""
    if not HA_TOKEN:
        return "HA_TOKEN not configured."
    headers = {"Authorization": f"Bearer {HA_TOKEN}", "Content-Type": "application/json"}
    try:
        async with httpx.AsyncClient(timeout=5.0) as client:
            if action == "get_states":
                if entity_id == "all":
                    r = await client.get(f"{HA_URL}/api/states", headers=headers)
                    states = r.json()
                    # Summarize lights/switches/sensors
                    lights = [s for s in states if s["entity_id"].startswith("light.")]
                    sensors = [s for s in states if s["entity_id"].startswith("sensor.")]
                    on_lights = [s["entity_id"].split(".")[1] for s in lights if s["state"] == "on"]
                    return (
                        f"{len(lights)} lights ({len(on_lights)} on: {', '.join(on_lights[:5])}), "
                        f"{len(sensors)} sensors"
                    )
                else:
                    r = await client.get(f"{HA_URL}/api/states/{entity_id}", headers=headers)
                    s = r.json()
                    return f"{entity_id}: {s.get('state')} (attrs: {list(s.get('attributes', {}).keys())[:3]})"
            else:
                service = action.replace("turn_", "").replace("toggle", "toggle")
                domain = entity_id.split(".")[0]
                r = await client.post(
                    f"{HA_URL}/api/services/{domain}/{action}",
                    headers=headers,
                    json={"entity_id": entity_id},
                )
                return f"{action} {entity_id}: {'done' if r.status_code in (200, 201) else r.text}"
    except Exception as e:
        return f"HA error: {e}"


@llm.function_tool
async def search_web(
    query: Annotated[str, "Search query"],
) -> str:
    """Search the web using the local SearXNG instance."""
    try:
        async with httpx.AsyncClient(timeout=8.0) as client:
            r = await client.get(
                f"{SEARXNG_URL}/search",
                params={"q": query, "format": "json", "categories": "general"},
            )
            r.raise_for_status()
            data = r.json()
        results = data.get("results", [])[:4]
        if not results:
            return "No results found."
        parts = []
        for res in results:
            title = res.get("title", "")
            content = res.get("content", "")[:150]
            parts.append(f"• {title}: {content}")
        return "\n".join(parts)
    except Exception as e:
        return f"Search failed: {e}"


ALL_TOOLS = [infra_status, container_action, check_logs, read_alerts, home_assistant, search_web]