322 lines
10 KiB
Python
322 lines
10 KiB
Python
import json
|
|
import os
|
|
import smtplib
|
|
import subprocess
|
|
from email.message import EmailMessage
|
|
from pathlib import Path
|
|
from typing import Iterable, TypedDict, cast
|
|
|
|
import requests
|
|
from dotenv import dotenv_values, load_dotenv
|
|
from openai import OpenAI
|
|
from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolUnionParam
|
|
|
|
|
|
class NTFY_DOTENV(TypedDict):
|
|
TOKEN: str
|
|
|
|
|
|
class ToolCallController:
|
|
def __init__(self, max_tool_calls=10):
|
|
self.max_tool_calls = max_tool_calls
|
|
self.tool_call_count = 0
|
|
|
|
def is_tool_call_allowed(self):
|
|
return self.tool_call_count < self.max_tool_calls
|
|
|
|
def increment(self):
|
|
self.tool_call_count += 1
|
|
|
|
def reset(self):
|
|
self.tool_call_count = 0
|
|
|
|
|
|
# Register tools
|
|
tools: Iterable[ChatCompletionToolUnionParam] = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "list_btrfs_pools",
|
|
"description": "List the btrfs pools on this system. Returns a list like ['/btrfs/pool0', '/btrfs/backup0', etc.]",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"server_id": {"type": "string", "enum": ["all"]}},
|
|
"required": ["server_id"],
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "btrfs_device_stats",
|
|
"description": "Runs `btrfs device stats` on a given pool.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"pool_path": {
|
|
"type": "string",
|
|
}
|
|
},
|
|
"required": ["pool_path"],
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "btrfs_filesystem_show",
|
|
"description": "Runs `btrfs filesystem show` on a given pool.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"pool_path": {
|
|
"type": "string",
|
|
}
|
|
},
|
|
"required": ["pool_path"],
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "btrfs_scrub_start",
|
|
"description": "Runs `btrfs scrub start` on a given pool.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"pool_path": {
|
|
"type": "string",
|
|
}
|
|
},
|
|
"required": ["pool_path"],
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "lsblk_filesystems",
|
|
"description": (
|
|
"Lists current filesystems on this server. Can be used to translate a "
|
|
"luks ID to a device. Example: luks-pool0-1c9a755a-2a81-4a2f-9bd7-6b6d7caaa523 "
|
|
"translates to /dev/sdb because /dev/sdb is listed below that filesystem. "
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"server_id": {"type": "string", "enum": ["all"]}},
|
|
"required": ["server_id"],
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "smartctl",
|
|
"description": (
|
|
"Runs `smartctl -a` on a given device. `device_path` should be "
|
|
"a device at /dev. Example:`/dev/sda` or `/dev/nvme0n1`."
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"device_path": {
|
|
"type": "string",
|
|
}
|
|
},
|
|
"required": ["device_path"],
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "alert_user",
|
|
"description": (
|
|
"Sends the provided message to the user via a messaging service."
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"message": {
|
|
"type": "string",
|
|
}
|
|
},
|
|
"required": ["message"],
|
|
},
|
|
},
|
|
},
|
|
]
|
|
|
|
|
|
def list_btrfs_pools(server_id: str) -> list[str]:
|
|
btrfs_path = Path("/btrfs")
|
|
pools = btrfs_path.iterdir()
|
|
list_pools = list(map(lambda item: str(item), pools))
|
|
return list_pools
|
|
|
|
|
|
def btrfs_device_stats(pool_path: str) -> str:
|
|
command_result = subprocess.run(
|
|
["btrfs", "device", "stats", pool_path], capture_output=True
|
|
)
|
|
stdout = command_result.stdout.decode()
|
|
stderr = command_result.stderr.decode()
|
|
output = f"**STDOUT**\n{stdout}\n\n**STDERR**{stderr}"
|
|
return output
|
|
|
|
|
|
def btrfs_filesystem_show(pool_path: str) -> str:
|
|
command_result = subprocess.run(
|
|
["btrfs", "filesystem", "show", pool_path], capture_output=True
|
|
)
|
|
stdout = command_result.stdout.decode()
|
|
stderr = command_result.stderr.decode()
|
|
output = f"**STDOUT**\n{stdout}\n\n**STDERR**{stderr}"
|
|
return output
|
|
|
|
|
|
def btrfs_scrub_start(pool_path: str) -> str:
|
|
command_result = subprocess.run(
|
|
["btrfs", "scrub", "start", pool_path], capture_output=True
|
|
)
|
|
stdout = command_result.stdout.decode()
|
|
stderr = command_result.stderr.decode()
|
|
output = f"**STDOUT**\n{stdout}\n\n**STDERR**{stderr}"
|
|
return output
|
|
|
|
|
|
def lsblk_filesystems(server_id: str) -> str:
|
|
command_result = subprocess.run(["lsblk", "-fs"], capture_output=True)
|
|
stdout = command_result.stdout.decode()
|
|
stderr = command_result.stderr.decode()
|
|
output = f"**STDOUT**\n{stdout}\n\n**STDERR**{stderr}"
|
|
return output
|
|
|
|
|
|
def smartctl(device_path: str) -> str:
|
|
command_result = subprocess.run(
|
|
["smartctl", "-a", device_path], capture_output=True
|
|
)
|
|
stdout = command_result.stdout.decode()
|
|
stderr = command_result.stderr.decode()
|
|
output = f"**STDOUT**\n{stdout}\n\n**STDERR**{stderr}"
|
|
return output
|
|
|
|
|
|
def load_ntfy_creds() -> NTFY_DOTENV:
|
|
ses_dotenv_location = Path(os.getenv("HOME", "/root"), ".env/aws_ses")
|
|
print(f"Loading env from {ses_dotenv_location}")
|
|
raw_values = dotenv_values(ses_dotenv_location)
|
|
if raw_values:
|
|
ntfy_config = cast(NTFY_DOTENV, raw_values)
|
|
# print(f"AWS SES Credentials loaded: {aws_ses_config}")
|
|
return ntfy_config
|
|
print("No ntfy credentials supplied. Exiting.")
|
|
exit(1)
|
|
|
|
|
|
def alert_user(message: str) -> str:
|
|
ntfy_creds = load_ntfy_creds()
|
|
|
|
requests.post(
|
|
"https://ntfy.reeseapps.com/servers",
|
|
data=message,
|
|
headers={
|
|
"Authorization": f"Bearer {ntfy_creds['TOKEN']}",
|
|
"Title": "Agent Disk Report",
|
|
"Priority": "default",
|
|
},
|
|
)
|
|
|
|
return "Sent."
|
|
|
|
|
|
def execute_tool(tool_name, arguments):
|
|
if tool_name == "list_btrfs_pools":
|
|
return list_btrfs_pools(**arguments)
|
|
elif tool_name == "btrfs_device_stats":
|
|
return btrfs_device_stats(**arguments)
|
|
elif tool_name == "btrfs_filesystem_show":
|
|
return btrfs_filesystem_show(**arguments)
|
|
elif tool_name == "btrfs_scrub_start":
|
|
return btrfs_scrub_start(**arguments)
|
|
elif tool_name == "lsblk_filesystems":
|
|
return lsblk_filesystems(**arguments)
|
|
elif tool_name == "smartctl":
|
|
return smartctl(**arguments)
|
|
elif tool_name == "alert_user":
|
|
return alert_user(**arguments)
|
|
raise ValueError(f"Unknown tool: {tool_name}")
|
|
|
|
|
|
def run_conversation(user_message: str, max_tool_calls=10):
|
|
print("Processing initial message")
|
|
controller = ToolCallController(max_tool_calls=max_tool_calls)
|
|
messages: Iterable[ChatCompletionMessageParam] = [
|
|
{
|
|
"role": "system",
|
|
"content": "You are a system administrator with access to a variety of administrator tools.",
|
|
}
|
|
]
|
|
|
|
messages.append({"role": "user", "content": user_message})
|
|
|
|
while True:
|
|
if not controller.is_tool_call_allowed():
|
|
messages.append(
|
|
{
|
|
"role": "user",
|
|
"content": "You've reached the maximum number of tool calls. Please summarize based on available information.",
|
|
}
|
|
)
|
|
break
|
|
|
|
response = client.chat.completions.create(
|
|
model="qwen3.5-35b-a3b", messages=messages, tools=tools, tool_choice="auto"
|
|
)
|
|
|
|
message = response.choices[0].message
|
|
print(f"\n{message}\n")
|
|
messages.append(message)
|
|
|
|
if message.tool_calls:
|
|
for tool_call in message.tool_calls:
|
|
controller.increment()
|
|
tool_name = tool_call.function.name
|
|
arguments = json.loads(tool_call.function.arguments)
|
|
print(f"Attempting to call {tool_name} with arguments {arguments}...")
|
|
|
|
# Give the user a chance to stop a problem before it starts
|
|
# keep_going = input("Continue? (Y/n) ")
|
|
# if keep_going.lower() == "n":
|
|
# exit(1)
|
|
|
|
result = execute_tool(tool_name, arguments)
|
|
|
|
messages.append(
|
|
{
|
|
"role": "tool",
|
|
"tool_call_id": tool_call.id,
|
|
"content": str(result),
|
|
}
|
|
)
|
|
else:
|
|
break
|
|
|
|
|
|
if __name__ == "__main__":
|
|
load_dotenv()
|
|
api_key = os.getenv("OPENAI_API_KEY", "")
|
|
client = OpenAI(base_url="https://aipi.reeseapps.com", api_key=api_key)
|
|
# Example usage
|
|
print(
|
|
run_conversation(
|
|
"Check the btrfs pools on this system. Take the appropriate action if any pools aren't "
|
|
"healthy. Don't run scrubs unless necessary. Also check the btrfs pool space and report "
|
|
"if any are getting full. At the very end of performing your checks send a single, "
|
|
"concise message to the user via the alert_user tool explaining what you did and "
|
|
"what concerns you might have."
|
|
)
|
|
)
|