initial working system report
This commit is contained in:
332
active/vibe_system_report/agent_report.py
Normal file
332
active/vibe_system_report/agent_report.py
Normal file
@@ -0,0 +1,332 @@
|
||||
import json
|
||||
import os
|
||||
import smtplib
|
||||
import ssl
|
||||
import subprocess
|
||||
from email.message import EmailMessage
|
||||
from pathlib import Path
|
||||
from typing import Iterable, TypedDict, cast
|
||||
|
||||
from dotenv import dotenv_values
|
||||
from openai import OpenAI
|
||||
from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolUnionParam
|
||||
|
||||
|
||||
class AWS_SES_DOTENV(TypedDict):
|
||||
USER: str
|
||||
PASSWORD: str
|
||||
ENDPOINT: str
|
||||
TLS_PORT: str
|
||||
SENDER: str
|
||||
RECEIVER: 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_ses_creds() -> AWS_SES_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:
|
||||
aws_ses_config = cast(AWS_SES_DOTENV, raw_values)
|
||||
print(f"AWS SES Credentials loaded: {aws_ses_config}")
|
||||
return aws_ses_config
|
||||
print("No email credentials supplied. Exiting.")
|
||||
exit(1)
|
||||
|
||||
|
||||
def alert_user(message: str) -> str:
|
||||
ses_config = load_ses_creds()
|
||||
port = int(ses_config["TLS_PORT"])
|
||||
user = ses_config["USER"]
|
||||
password = ses_config["PASSWORD"]
|
||||
sender = ses_config["SENDER"]
|
||||
receiver = ses_config["RECEIVER"]
|
||||
|
||||
# Create a secure SSL context
|
||||
context = ssl.create_default_context()
|
||||
|
||||
msg = EmailMessage()
|
||||
msg["Subject"] = "Agent Disk Report"
|
||||
msg["From"] = sender
|
||||
msg["To"] = receiver
|
||||
msg.set_content(message)
|
||||
|
||||
with smtplib.SMTP_SSL(ses_config["ENDPOINT"], port, context=context) as server:
|
||||
server.login(user, password)
|
||||
result = server.send_message(msg)
|
||||
print(result)
|
||||
|
||||
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__":
|
||||
client = OpenAI(base_url="https://llama-cpp.reeselink.com", 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 explaining what you did and what concerns you might have."
|
||||
)
|
||||
)
|
||||
Reference in New Issue
Block a user