From 7f4b55964463d189f766256688b9d8caa6f04cda Mon Sep 17 00:00:00 2001
From: "Max P." <mail@0xMax42.io>
Date: Sun, 5 Oct 2025 20:33:27 +0200
Subject: [PATCH] feat(module): add support for llama-swap API integration

---
 src/pyvtt/libs/ollama.py | 113 +++++++++++++++++++++++++--------------
 1 file changed, 74 insertions(+), 39 deletions(-)

diff --git a/src/pyvtt/libs/ollama.py b/src/pyvtt/libs/ollama.py
index 061da5e..1b0f337 100644
--- a/src/pyvtt/libs/ollama.py
+++ b/src/pyvtt/libs/ollama.py
@@ -1,3 +1,4 @@
+import json
 import requests
 from typing import Union, List, Optional
 
@@ -8,61 +9,95 @@ from pyvtt.models.config import AppConfig, PresetConfig
 class OllamaClient:
     def __init__(self, config: AppConfig):
         """
-        Initialisiert den Ollama-Client mit der Basis-Konfiguration aus der globalen App-Konfiguration.
-
-        :param config: AppConfig-Instanz mit Host und Port für den Ollama-Server.
+        Initialisiert den API-Client (Ollama oder llama-swap) mit Basis-Konfiguration.
         """
-        self.base_url = config.ollama_url
+        self.base_url = config.ollama_url.rstrip("/")
         self.port = config.ollama_port
-        self.path = config.ollama_path
+        self.path = config.ollama_path or "/api/chat"
+        # Falls llama-swap (OpenAI-API-Kompatibel), verwende den OpenAI-Pfad
+        if "v1" in self.path or "completions" in self.path:
+            self.is_llama_swap = True
+        else:
+            self.is_llama_swap = False
 
-    def send_chat(
-        self,
-        user_message: str,
-        config: PresetConfig,
-    ) -> str:
+    def send_chat(self, user_message: str, config: PresetConfig) -> str:
         """
-        Sendet eine Chat-Anfrage an den Ollama-Server basierend auf der spezifischen Preset-Konfiguration.
-
-        :param user_message: Der vom Nutzer erzeugte Eingabetext (z. B. Transkript).
-        :param config: PresetConfig-Instanz mit modell-, prompt- und kontextbezogenen Parametern.
-        :return: Der von Ollama zurückgegebene, formatierte Antworttext, die user_message
-                 unverändert zurückgibt, wenn Ollama deaktiviert ist oder none bei einem Fehler.
+        Sendet eine Chat-Anfrage an Ollama oder llama-swap.
         """
         if config.ollama and config.ollama.lower() == "disable":
             print("[OllamaClient] Ollama ist im Preset deaktiviert.")
-            print("[OllamaClient] Gebe die Eingabe unverändert zurück.")
             return user_message
 
-        # Prompt als String aufbereiten – Liste wird zu Zeilen verbunden
-        if isinstance(config.ollama_prompt, list):
-            prompt_str = "\n".join(config.ollama_prompt)
-        else:
-            prompt_str = config.ollama_prompt
+        # Prompt aufbereiten
+        prompt_str = (
+            "\n".join(config.ollama_prompt)
+            if isinstance(config.ollama_prompt, list)
+            else str(config.ollama_prompt)
+        )
 
-        # Payload für die API-Anfrage vorbereiten
-        payload = {
-            "model": config.ollama_model,
-            "messages": [
-                {"role": "system", "content": prompt_str},
-                {"role": "user", "content": user_message}
-            ],
-            "options": {
-                "num_ctx": config.ollama_context,
-            } if config.ollama_context else {},
-            "stream": False
-        }
+        # === Payload vorbereiten ===
+        if self.is_llama_swap:
+            # OpenAI-/llama-swap-kompatibles Format
+            payload = {
+                "model": config.ollama_model,
+                "messages": [
+                    {"role": "system", "content": prompt_str},
+                    {"role": "user", "content": user_message},
+                ],
+                "stream": False,
+            }
+            # Kontextgröße optional hinzufügen
+            if config.ollama_context:
+                payload["num_ctx"] = config.ollama_context
+        else:
+            # Klassisches Ollama-Format
+            payload = {
+                "model": config.ollama_model,
+                "messages": [
+                    {"role": "system", "content": prompt_str},
+                    {"role": "user", "content": user_message},
+                ],
+                "options": (
+                    {"num_ctx": config.ollama_context}
+                    if config.ollama_context
+                    else {}
+                ),
+                "stream": False,
+            }
 
         endpoint = f"{self.base_url}:{self.port}{self.path}"
 
-        # Anfrage an Ollama senden und Antwort extrahieren
+        # === Anfrage senden ===
         try:
-            response = requests.post(endpoint, json=payload)
+            headers = {"Content-Type": "application/json"}
+            if self.is_llama_swap:
+                headers["Authorization"] = "Bearer no-key"
+
+            response = requests.post(endpoint, headers=headers, data=json.dumps(payload))
             response.raise_for_status()
+
             json_response = response.json()
-            content = json_response.get("message", {}).get("content", "").strip()
+
+            # === Antwort extrahieren ===
+            if self.is_llama_swap:
+                # OpenAI-kompatible Struktur
+                content = (
+                    json_response.get("choices", [{}])[0]
+                    .get("message", {})
+                    .get("content", "")
+                    .strip()
+                )
+            else:
+                # Ollama-eigene Struktur
+                content = (
+                    json_response.get("message", {})
+                    .get("content", "")
+                    .strip()
+                )
+
             return "\n".join(line.strip() for line in content.splitlines())
+
         except requests.exceptions.RequestException as e:
             print(f"[OllamaClient] HTTP-Fehler: {e}")
-            notify("Fehler", "Ein Fehler bei der Kommunikation mit 'Ollama' ist aufgetreten!")
-            return ""
+            notify("Fehler", "Kommunikationsfehler mit Ollama / llama-swap!")
+            return ""
\ No newline at end of file