checkpoint commit

active/vibe_agent/main.py

@@ -0,0 +1,178 @@
+import json
+import subprocess
+from typing import Iterable
+
+from openai import OpenAI
+from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolUnionParam
+
+client = OpenAI(base_url="https://llama-cpp.reeselink.com", api_key="")
+
+
+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_servers",
+            "description": "Lists the available servers to perform operations on",
+            "parameters": {
+                "type": "object",
+                "properties": {"server_id": {"type": "string", "enum": ["all"]}},
+                "required": ["server_id"],
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "check_updates",
+            "description": "Check if a given server needs updated.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "server_id": {
+                        "type": "string",
+                    }
+                },
+                "required": ["server_id"],
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "perform_updates",
+            "description": "Update a given server to the latest package versions. Does not reboot automatically.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "server_id": {
+                        "type": "string",
+                    }
+                },
+                "required": ["server_id"],
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "reboot",
+            "description": "Reboot a given server. Waits for server to be responsive again.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "server_id": {
+                        "type": "string",
+                    }
+                },
+                "required": ["server_id"],
+            },
+        },
+    },
+]
+
+
+def list_servers() -> str:
+    return ",".join(["ignite"])
+
+
+def check_updates(server_id: str):
+    command_result = subprocess.run(
+        ["ssh", server_id, "dnf", "check-update"], capture_output=True
+    )
+    output = command_result.stdout.decode()
+    return output
+
+
+def perform_updates(server_id: str):
+    return f"Successfully updates {server_id}. Reboot required."
+
+
+def reboot(server_id: str):
+    return f"Rebooted {server_id} successfully."
+
+
+def execute_tool(tool_name, arguments):
+    if tool_name == "check_updates":
+        return check_updates(**arguments)
+    elif tool_name == "list_servers":
+        return list_servers()
+    elif tool_name == "perform_updates":
+        return perform_updates(**arguments)
+    elif tool_name == "reboot":
+        return reboot(**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
+        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}...")
+                result = execute_tool(tool_name, arguments)
+
+                messages.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": tool_call.id,
+                        "content": str(result),
+                    }
+                )
+        else:
+            break
+
+    return messages[-1].content
+
+
+# Example usage
+print(
+    run_conversation(
+        "Can you check if any of my servers need updates and update them? If they need rebooted, reboot them."
+    )
+)