diff --git a/main.py b/main.py index 483ca7d..c5f3ebf 100644 --- a/main.py +++ b/main.py @@ -4,6 +4,9 @@ import fcntl import subprocess import sqlite3 import asyncio +import openai +import google.generativeai as genai +import json from fastapi import FastAPI, WebSocket, BackgroundTasks, Request, Form, WebSocketDisconnect from fastapi.responses import RedirectResponse from fastapi.templating import Jinja2Templates @@ -17,6 +20,46 @@ templates = Jinja2Templates(directory="templates") SSH_KEY = os.path.expanduser("~/.ssh/id_rsa") DB_PATH = "cluster.db" +# --- KI KONFIGURATION --- +AI_PROVIDER = "ollama" # "openai", "google" oder "ollama" +OPENAI_API_KEY = "dein-key" +GOOGLE_API_KEY = "dein-key" +OLLAMA_BASE_URL = "http://x.x.x.x:11434/v1" # IP deines Ollama-Servers +# System Prompt: Sagt der KI, wer sie ist und was sie kann +SYSTEM_PROMPT = """Du bist der Pi-Orchestrator KI-Assistent. +Deine Aufgabe ist es, dem Nutzer zu helfen, Raspberry Pis zu verwalten und Docker-Container zu steuern. +Du kannst Linux-Befehle generieren. Antworte präzise und hilfsbereit.""" + +# --- KI FUNKTIONEN --- + +async def get_ai_response(user_input): + if AI_PROVIDER == "openai": + client = openai.OpenAI(api_key=OPENAI_API_KEY) + response = client.chat.completions.create( + model="gpt-4o", + messages=[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_input}] + ) + return response.choices[0].message.content + + elif AI_PROVIDER == "ollama": + # Ollama nutzt das OpenAI-Format, braucht aber keinen Key + client = openai.OpenAI(base_url=OLLAMA_BASE_URL, api_key="ollama") + response = client.chat.completions.create( + model="llama3", # Oder dein bevorzugtes Modell + messages=[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_input}] + ) + return response.choices[0].message.content + + elif AI_PROVIDER == "google": + genai.configure(api_key=GOOGLE_API_KEY) + model = genai.GenerativeModel('gemini-1.5-flash') + # Gemini braucht einen etwas anderen Aufbau für System Prompts + response = model.generate_content(f"{SYSTEM_PROMPT}\n\nNutzer: {user_input}") + return response.text + + return "Fehler: Kein KI-Provider konfiguriert." + + # --- DATENBANK INITIALISIERUNG --- def init_db(): conn = sqlite3.connect(DB_PATH) @@ -161,27 +204,19 @@ async def terminal_websocket(websocket: WebSocket, ip: str): os.close(master_fd) os.close(slave_fd) +# --- WEBSOCKET CHAT UPDATE --- + @app.websocket("/ws/chat") -async def chat_endpoint(websocket: WebSocket): +async def websocket_chat(websocket: WebSocket): await websocket.accept() - while True: - user_msg = await websocket.receive_text() - user_msg_lower = user_msg.lower() - - if "installiere docker" in user_msg_lower: - conn = get_db() - nodes = conn.execute('SELECT * FROM nodes').fetchall() - conn.close() - target = next((n for n in nodes if n['name'].lower() in user_msg_lower or n['ip'] in user_msg_lower), None) - - if target: - await websocket.send_text(f"🤖 OK. Starte Docker-Installation auf {target['name']}...") - cmd = "curl -sSL https://get.docker.com | sh && sudo usermod -aG docker " + target['user'] - asyncio.create_task(run_remote_task(target['ip'], target['user'], cmd, "Docker Installation")) - else: - await websocket.send_text("🤖 Node nicht gefunden. Welchen meinst du?") - else: - await websocket.send_text(f"🤖 Empfangen: {user_msg}. Soll ich etwas installieren?") + try: + while True: + data = await websocket.receive_text() + # Hole Antwort von der gewählten KI + ai_msg = await get_ai_response(data) + await websocket.send_text(ai_msg) + except Exception as e: + print(f"Chat Error: {e}") # --- Status in DB aktualisieren Helper --- def update_node_status(ip, new_status):