231 lines
9.2 KiB
Python
231 lines
9.2 KiB
Python
import os
|
|
import logging
|
|
from typing import TypedDict
|
|
from langchain_core.messages import SystemMessage, HumanMessage
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.graph import StateGraph, END
|
|
# AsyncPostgresSaver imported lazily below (only when conn_string is provided)
|
|
|
|
from ops_agent.tools import (
|
|
k8s_get_status, k8s_get_logs, k8s_restart_pod, k8s_scale_deployment,
|
|
proxmox_get_status, proxmox_get_disk_space, proxmox_reboot_container, proxmox_clean_disk
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class OpsState(TypedDict):
|
|
target_type: str # 'kubernetes' or 'proxmox'
|
|
target_id: str # e.g., 'default/web-pod' or VMID '101'
|
|
status_info: str # status output
|
|
logs: str # logs output
|
|
diagnosis: str # LLM diagnosis of issues
|
|
proposed_action: str # Proposed action: 'restart_pod', 'scale_deployment', 'reboot_container', 'clean_disk', 'none'
|
|
proposed_action_details: str # E.g., pod name, deployment name, or container ID
|
|
is_approved: bool # Has user approved it?
|
|
action_result: str # Outcome of the action execution
|
|
|
|
def get_llm():
|
|
"""Initializes LLM using environment variables, matching the Gumbo setup."""
|
|
api_key = os.environ.get("LITELLM_API_KEY") or os.environ.get("OPENAI_API_KEY", "mock-api-key")
|
|
base_url = os.environ.get("LITELLM_BASE_URL") or os.environ.get("OPENAI_BASE_URL", "http://localhost:4000")
|
|
model = os.environ.get("OPS_LLM_MODEL") or os.environ.get("GUMBO_LLM_MODEL", "ollama-qwen")
|
|
|
|
logger.info(f"Initializing ChatOpenAI with model={model}, base_url={base_url}")
|
|
return ChatOpenAI(
|
|
model=model,
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
temperature=0.1
|
|
)
|
|
|
|
async def _fetch_status_and_logs(state: OpsState) -> dict:
|
|
target_type = state.get("target_type", "").lower()
|
|
target_id = state.get("target_id", "")
|
|
|
|
status_info = ""
|
|
logs = ""
|
|
|
|
if target_type == "kubernetes":
|
|
# target_id pattern: "namespace/pod-name" or just "pod-name" (assumes default namespace)
|
|
parts = target_id.split("/")
|
|
if len(parts) == 2:
|
|
ns, pod = parts[0], parts[1]
|
|
else:
|
|
ns, pod = "default", parts[0]
|
|
|
|
status_info = k8s_get_status(ns, pod)
|
|
logs = k8s_get_logs(ns, pod)
|
|
|
|
elif target_type == "proxmox":
|
|
# target_id is LXC VMID (e.g. "101")
|
|
status_info = proxmox_get_status(target_id)
|
|
logs = proxmox_get_disk_space(target_id)
|
|
|
|
else:
|
|
status_info = f"Unknown target type: {target_type}"
|
|
logs = "No logs available."
|
|
|
|
return {"status_info": status_info, "logs": logs}
|
|
|
|
async def _diagnose_issue(state: OpsState) -> dict:
|
|
llm = get_llm()
|
|
system_msg = SystemMessage(
|
|
content=(
|
|
"You are OpsAgent, a senior DevOps engineer and Proxmox/Kubernetes site reliability assistant. "
|
|
"Analyze the provided status info and logs, identify the root cause of any problems, and return "
|
|
"a concise Markdown diagnosis."
|
|
)
|
|
)
|
|
user_content = (
|
|
f"Target Type: {state['target_type']}\n"
|
|
f"Target ID: {state['target_id']}\n\n"
|
|
f"=== STATUS INFO ===\n{state['status_info']}\n\n"
|
|
f"=== LOGS / RESOURCE USAGE ===\n{state['logs']}\n"
|
|
)
|
|
|
|
try:
|
|
resp = await llm.ainvoke([system_msg, HumanMessage(content=user_content)])
|
|
diagnosis = resp.content if isinstance(resp.content, str) else str(resp.content)
|
|
except Exception as e:
|
|
logger.error(f"LLM call failed in diagnosis: {e}")
|
|
diagnosis = f"LLM error: Failed to diagnose log. Raw error: {str(e)}"
|
|
|
|
return {"diagnosis": diagnosis}
|
|
|
|
async def _propose_remediation(state: OpsState) -> dict:
|
|
llm = get_llm()
|
|
system_msg = SystemMessage(
|
|
content=(
|
|
"You are OpsAgent. Based on the diagnosis, propose a specific remediation action from the following list:\n"
|
|
"- restart_pod (Kubernetes pods crash/hangs)\n"
|
|
"- scale_deployment (K8s capacity issues)\n"
|
|
"- reboot_container (Proxmox container stopped/unresponsive)\n"
|
|
"- clean_disk (Proxmox/LXC container running out of disk space: Use% >= 90%)\n"
|
|
"- none (No actions required, target is running fine)\n\n"
|
|
"Respond ONLY with a JSON object: "
|
|
'{"action": "restart_pod"|"scale_deployment"|"reboot_container"|"clean_disk"|"none", '
|
|
'"details": "brief reason or extra target parameter"}'
|
|
)
|
|
)
|
|
user_content = f"Diagnosis:\n{state['diagnosis']}"
|
|
|
|
action = "none"
|
|
details = "No action proposed."
|
|
|
|
try:
|
|
resp = await llm.ainvoke([system_msg, HumanMessage(content=user_content)])
|
|
content = resp.content if isinstance(resp.content, str) else str(resp.content)
|
|
# Parse JSON
|
|
import json
|
|
# Strip potential markdown fences
|
|
clean_content = content.replace("```json", "").replace("```", "").strip()
|
|
data = json.loads(clean_content)
|
|
action = data.get("action", "none")
|
|
details = data.get("details", "")
|
|
|
|
# Dispatch WhatsApp Alert to phone if an anomaly/remediation is flagged
|
|
if action != "none":
|
|
try:
|
|
from ops_agent.tools import send_whatsapp_alert
|
|
alert_msg = (
|
|
f"⚠️ *OPSAGENT THREAT WARNING* ⚠️\n\n"
|
|
f"*Target:* {state['target_type'].upper()} ({state['target_id']})\n"
|
|
f"*Diagnosis:* {details}\n"
|
|
f"*Proposed Action:* {action.replace('_', ' ').upper()}\n\n"
|
|
f"Login to the Agentic OS Dashboard to authorize or reject this mitigation."
|
|
)
|
|
send_whatsapp_alert(alert_msg)
|
|
except Exception as ex:
|
|
logger.error(f"Could not dispatch WhatsApp alert: {ex}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to propose action via LLM: {e}")
|
|
|
|
return {"proposed_action": action, "proposed_action_details": details}
|
|
|
|
async def _execute_remediation(state: OpsState) -> dict:
|
|
if not state.get("is_approved", False):
|
|
return {"action_result": "Remediation action skipped: Approval was not provided."}
|
|
|
|
action = state.get("proposed_action", "none")
|
|
target_type = state.get("target_type", "")
|
|
target_id = state.get("target_id", "")
|
|
|
|
result = "No action executed."
|
|
|
|
try:
|
|
if target_type == "kubernetes":
|
|
parts = target_id.split("/")
|
|
ns = parts[0] if len(parts) == 2 else "default"
|
|
pod = parts[1] if len(parts) == 2 else parts[0]
|
|
|
|
if action == "restart_pod":
|
|
result = k8s_restart_pod(ns, pod)
|
|
elif action == "scale_deployment":
|
|
# Default scale to 2 for demo purposes
|
|
result = k8s_scale_deployment(ns, pod, 2)
|
|
else:
|
|
result = f"Action '{action}' is not supported for Kubernetes target."
|
|
|
|
elif target_type == "proxmox":
|
|
if action == "reboot_container":
|
|
result = proxmox_reboot_container(target_id)
|
|
elif action == "clean_disk":
|
|
result = proxmox_clean_disk(target_id)
|
|
else:
|
|
result = f"Action '{action}' is not supported for Proxmox container VMID {target_id}."
|
|
else:
|
|
result = f"Unknown target type: {target_type}."
|
|
|
|
except Exception as e:
|
|
result = f"Error during execution: {str(e)}"
|
|
|
|
return {"action_result": result}
|
|
|
|
|
|
def build_graph_builder() -> StateGraph:
|
|
builder = StateGraph(OpsState)
|
|
builder.add_node("fetch_status", _fetch_status_and_logs)
|
|
builder.add_node("diagnose", _diagnose_issue)
|
|
builder.add_node("propose_action", _propose_remediation)
|
|
builder.add_node("execute_action", _execute_remediation)
|
|
|
|
builder.set_entry_point("fetch_status")
|
|
builder.add_edge("fetch_status", "diagnose")
|
|
builder.add_edge("diagnose", "propose_action")
|
|
builder.add_edge("propose_action", "execute_action")
|
|
builder.add_edge("execute_action", END)
|
|
|
|
return builder
|
|
|
|
async def run_ops_graph(target_type: str, target_id: str, is_approved: bool = False, conn_string: str = None) -> OpsState:
|
|
builder = build_graph_builder()
|
|
|
|
initial_state = {
|
|
"target_type": target_type,
|
|
"target_id": target_id,
|
|
"status_info": "",
|
|
"logs": "",
|
|
"diagnosis": "",
|
|
"proposed_action": "none",
|
|
"proposed_action_details": "",
|
|
"is_approved": is_approved,
|
|
"action_result": ""
|
|
}
|
|
|
|
if conn_string:
|
|
try:
|
|
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
|
except ImportError:
|
|
from langgraph_checkpoint_postgres.aio import AsyncPostgresSaver
|
|
async with AsyncPostgresSaver.from_conn_string(conn_string) as checkpointer:
|
|
await checkpointer.setup()
|
|
graph = builder.compile(checkpointer=checkpointer)
|
|
return await graph.ainvoke(
|
|
initial_state,
|
|
config={"configurable": {"thread_id": f"ops-{target_type}-{target_id}"}},
|
|
)
|
|
else:
|
|
graph = builder.compile()
|
|
return await graph.ainvoke(initial_state)
|