services.py

import os
import sys
import shutil
import subprocess
import requests
import threading
import time
import re
from PyQt6.QtCore import QObject, pyqtSignal
from config import get_base_path, MODELS_DIR

def clean_ansi(text): 
    return re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])').sub('', text)

class ServiceManager:
    def __init__(self):
        self.process = None
        base = get_base_path()
        self.ollama_path = os.path.join(base, "bin", "ollama.exe")
        if not os.path.exists(self.ollama_path):
            if not getattr(sys, 'frozen', False): 
                self.ollama_path = shutil.which("ollama")
        self.env = os.environ.copy()
        self.env["OLLAMA_MODELS"] = MODELS_DIR
        
    def start_service(self):
        if not self.ollama_path: 
            return False, "No Ollama"
        try:
            if requests.get("http://localhost:11434", timeout=0.2).status_code == 200: 
                return True, "Attached"
        except: 
            pass
        try:
            si = subprocess.STARTUPINFO()
            si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            self.process = subprocess.Popen(
                [self.ollama_path, "serve"], 
                env=self.env, 
                stdout=subprocess.PIPE, 
                stderr=subprocess.PIPE, 
                startupinfo=si, 
                creationflags=subprocess.CREATE_NO_WINDOW
            )
            time.sleep(1)
            return True, "Started Local"
        except Exception as e: 
            return False, str(e)
            
    def stop_service(self):
        if self.process: 
            # ИСПОЛЬЗУЕМ KILL ВМЕСТО TERMINATE ДЛЯ МГНОВЕННОГО ЗАКРЫТИЯ
            self.process.kill()
            
    def list_local_models(self):
        try:
            si = subprocess.STARTUPINFO()
            si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            res = subprocess.run(
                [self.ollama_path, "list"], 
                env=self.env, 
                capture_output=True, 
                text=True, 
                encoding='utf-8', 
                startupinfo=si, 
                creationflags=subprocess.CREATE_NO_WINDOW
            )
            if res.returncode == 0: 
                return [l.split()[0] for l in res.stdout.strip().split('\n')[1:] if l.strip()]
        except: 
            pass
        return []
        
    def delete_model(self, m):
        try:
            si = subprocess.STARTUPINFO()
            si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            subprocess.run(
                [self.ollama_path, "rm", m], 
                env=self.env, 
                startupinfo=si, 
                creationflags=subprocess.CREATE_NO_WINDOW
            )
            return True
        except: 
            return False
            
    def get_gpu_info(self):
        try:
            si = subprocess.STARTUPINFO()
            si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            res = subprocess.run(
                ['nvidia-smi', '--query-gpu=name,memory.total', '--format=csv,noheader'], 
                capture_output=True, 
                text=True, 
                startupinfo=si, 
                creationflags=subprocess.CREATE_NO_WINDOW
            )
            if res.returncode == 0:
                name, mem = res.stdout.strip().split(',')
                return name, float(mem.lower().replace('mib','').strip())/1024.0
        except: 
            pass
        return "Unknown GPU", 0

class ModelInstaller(QObject):
    progress_signal = pyqtSignal(int, str)
    finished_signal = pyqtSignal(bool)
    
    def __init__(self, model_name, svc): 
        super().__init__()
        self.model_name = model_name
        self.svc = svc
        self.running = True
        self.proc = None
        
    def run(self): 
        threading.Thread(target=self._pull_thread, daemon=True).start()
        
    def stop(self):
        self.running = False
        if self.proc: 
            self.proc.kill()
            
    def _pull_thread(self):
        path = self.svc.ollama_path if self.svc.ollama_path else "ollama"
        try:
            si = subprocess.STARTUPINFO()
            si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            self.proc = subprocess.Popen(
                [path, "pull", self.model_name], 
                env=self.svc.env, 
                stdout=subprocess.PIPE, 
                stderr=subprocess.STDOUT, 
                text=True, 
                encoding='utf-8', 
                errors='replace', 
                startupinfo=si, 
                creationflags=subprocess.CREATE_NO_WINDOW
            )
            for line in self.proc.stdout:
                if not self.running: break
                clean = clean_ansi(line).strip()
                if clean:
                    if "pulling" in line: self.progress_signal.emit(-1, clean)
                    elif "verifying" in line: self.progress_signal.emit(99, "Verifying...")
            self.proc.wait()
            if self.running and self.proc.returncode == 0: 
                self.progress_signal.emit(100, "OK")
                self.finished_signal.emit(True)
            else: 
                self.progress_signal.emit(0, "Stopped")
                self.finished_signal.emit(False)
        except Exception as e: 
            self.progress_signal.emit(0, str(e))
            self.finished_signal.emit(False)