Skip to content

C
coderai
Commit df366b63 authored by Your Name's avatar Your Name
Browse files
Options

Remove duplicate class/function definitions from coderai 

Removed ~1500 lines of duplicate code that now exist in codai modules:
- ModelCapabilities, detect_model_capabilities (now in codai.models.capabilities)
- Cache functions (now in codai.models.cache)
- detect_available_backends, check_flash_attn_availability (now in codai.backends)
- ModelBackend abstract class (now in codai.backends.base)
- ModelManager, WhisperServerManager, MultiModelManager (now in codai.models.manager)
- QueueManager (now in codai.queue.manager)
- Utility functions (now in codai.models.utils)

The code now properly imports from codai modules instead of having inline duplicates.
parent 9299c34f
Show whitespace changes
Inline Side-by-side
Showing with 0 additions and 1545 deletions
coderai
View file @ df366b63
......@@ -48,249 +48,14 @@ queue_flags = {"model_1": False, "image_1": False, "audio_1": False, "tts_1": Fa
# Model Cache Directory
# =============================================================================
# =============================================================================
# Model Capability Detection
# =============================================================================
from typing import Set, Optional
from dataclasses import dataclass, field
@dataclass
class ModelCapabilities:
"""Represents what a model can do."""
text_generation: bool = False # LLM/chat completion
image_to_text: bool = False # Image understanding (captioning, VQA)
image_generation: bool = False # Text-to-image (Stable Diffusion)
speech_to_text: bool = False # Audio transcription
text_to_speech: bool = False # Speech synthesis
def __str__(self):
caps = []
if self.text_generation:
caps.append("text")
if self.image_to_text:
caps.append("image-to-text")
if self.image_generation:
caps.append("image")
if self.speech_to_text:
caps.append("speech-to-text")
if self.text_to_speech:
caps.append("text-to-speech")
return ", ".join(caps) if caps else "none"
def detect_model_capabilities(model_name: str) -> ModelCapabilities:
"""
Detect model capabilities based on model name/type.
This is a heuristic detection - actual capabilities may vary.
"""
caps = ModelCapabilities()
if not model_name:
return caps
name_lower = model_name.lower()
# Check for image generation models (Stable Diffusion, SDXL, etc.)
if any(x in name_lower for x in ['stable-diffusion', 'sd15', 'sdxl', 'sd-xl', 'turbo', 'playground']):
caps.image_generation = True
return caps # Usually SD models are dedicated
# Check for vision models (image-to-text)
if any(x in name_lower for x in ['vision', 'vl-', '-vl', 'llava', 'qwen2-vl', 'qwen-vl', 'phi-4-mini', 'pixtral', 'clip']):
caps.image_to_text = True
caps.text_generation = True # Vision models are also LLMs
return caps
# Check for TTS models
if any(x in name_lower for x in ['kokoro', 'tts', 'speech', 'voice']):
caps.text_to_speech = True
return caps
# Check for whisper models (speech-to-text)
if any(x in name_lower for x in ['whisper', 'faster-whisper', 'distil-whisper']):
caps.speech_to_text = True
return caps
# Check for GGUF models (typically text models)
if '.gguf' in name_lower or 'gguf' in name_lower:
caps.text_generation = True
return caps
# Default: assume text generation (most HF models are LLMs)
caps.text_generation = True
return caps
def get_model_cache_dir() -> str:
"""Get or create the model cache directory."""
# Use XDG_CACHE_HOME if set, otherwise use ~/.cache/coderai
cache_home = os.environ.get('XDG_CACHE_HOME', os.path.expanduser('~/.cache'))
cache_dir = os.path.join(cache_home, 'coderai', 'models')
pathlib.Path(cache_dir).mkdir(parents=True, exist_ok=True)
return cache_dir
def get_all_cache_dirs() -> dict:
"""Get all model cache directories."""
caches = {}
cache_home = os.environ.get('XDG_CACHE_HOME', os.path.expanduser('~/.cache'))
# Coderai GGUF cache
coderai_cache = os.path.join(cache_home, 'coderai', 'models')
if os.path.exists(coderai_cache):
caches['coderai'] = coderai_cache
# HuggingFace cache (for .safetensors, PyTorch models, etc.)
hf_cache = os.path.join(cache_home, 'huggingface')
if os.path.exists(hf_cache):
caches['huggingface'] = hf_cache
# Local diffusers cache (often stored locally by apps)
local_diffusers = os.path.expanduser('~/.cache/diffusers')
if os.path.exists(local_diffusers):
caches['diffusers'] = local_diffusers
return caches
def get_cached_model_path(url: str) -> Optional[str]:
"""Check if a model URL is already cached. Returns path if cached, None otherwise."""
cache_dir = get_model_cache_dir()
# Create a safe filename from the URL using SHA-256 hash
url_hash = hashlib.sha256(url.encode()).hexdigest()
url_path = url.split('?')[0] # Remove query params
filename = os.path.basename(url_path)
if not filename:
filename = "model.gguf"
# Match the format used in load_model: {hash}_{filename}
cached_filename = f"{url_hash}_{filename}"
cached_path = os.path.join(cache_dir, cached_filename)
if os.path.exists(cached_path):
print(f"Using cached model: {cached_path}")
return cached_path
return None
def is_huggingface_model_id(path: str) -> bool:
"""Check if the path is a Hugging Face model ID (e.g., 'Qwen/Qwen3-4B-Instruct-2507-Q3_K_S')."""
# Must contain / but not be a URL
return '/' in path and not path.startswith('http://') and not path.startswith('https://')
def download_huggingface_model(model_id: str, cache_dir: str, file_pattern: str = '.gguf') -> Optional[str]:
"""Download a model from Hugging Face by model ID. Returns cached path or None on failure."""
try:
from huggingface_hub import hf_hub_download, list_repo_files
print(f"Downloading from Hugging Face: {model_id}")
files = list_repo_files(model_id)
# Filter files by extension pattern
matching_files = [f for f in files if f.endswith(file_pattern)]
if not matching_files:
print(f"No {file_pattern} files found in {model_id}")
return None
# Download the first matching file
filename = matching_files[0]
print(f"Downloading: {filename}")
model_path = hf_hub_download(repo_id=model_id, filename=filename, cache_dir=cache_dir)
print(f"Downloaded model to: {model_path}")
return model_path
except Exception as e:
print(f"Error downloading from Hugging Face: {e}")
return None
def download_model(url: str, cache_dir: str) -> str:
"""Download a model from URL with progress reporting. Returns cached path."""
import requests
import hashlib
url_path = url.split('?')[0]
filename = os.path.basename(url_path)
# Determine file extension
if 'gguf' in url.lower():
ext = '.gguf'
elif 'bin' in url.lower():
ext = '.bin'
elif 'ggml' in url.lower():
ext = '.ggml'
else:
ext = '.bin'
if not filename.endswith(ext):
filename = f"model{ext}"
# Create safe filename in cache
url_hash = hashlib.sha256(url.encode()).hexdigest()
cached_filename = f"{url_hash}_{filename}"
model_path = os.path.join(cache_dir, cached_filename)
# Check if already cached
if os.path.exists(model_path):
print(f"Using cached model: {model_path}")
return model_path
# Download
print(f"Downloading model: {url}")
response = requests.get(url, stream=True)
response.raise_for_status()
total_size = int(response.headers.get('content-length', 0))
total_mb = total_size / (1024 * 1024) if total_size > 0 else 0
downloaded = 0
start_time = time.time()
with open(model_path, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192*1024):
if chunk:
f.write(chunk)
downloaded += len(chunk)
if total_size > 0:
percent = (downloaded / total_size) * 100
elapsed = time.time() - start_time
speed = downloaded / elapsed if elapsed > 0 else 0
speed_mb = speed / (1024 * 1024)
dl_mb = downloaded / (1024 * 1024)
print(f"Downloaded: {percent:.1f}% ({dl_mb:.1f}/{total_mb:.1f} MB) at {speed_mb:.1f} MB/s", end='\r')
print() # New line after progress
print(f"Downloaded and cached to: {model_path}")
if total_mb > 0:
print(f"File size: {total_mb:.1f} MB")
return model_path
# =============================================================================
# Backend Detection and Imports
# =============================================================================
def detect_available_backends():
"""Detect which backends are available."""
backends = {'cpu': True}
# Check for PyTorch/CUDA
try:
import torch
if torch.cuda.is_available():
backends['nvidia'] = True
except ImportError:
pass
# Check for llama-cpp-python (Vulkan)
try:
import llama_cpp
backends['vulkan'] = True
except ImportError:
pass
return backends
# =============================================================================
# Flash Attention Detection (for NVIDIA backend)
# =============================================================================
def check_flash_attn_availability() -> bool:
"""Check if flash-attn is installed and available."""
try:
import flash_attn
return True
except ImportError:
return False
# =============================================================================
# Pydantic Models for API
# =============================================================================
......@@ -910,42 +675,6 @@ def format_tools_for_prompt(tools: List[Tool], messages: List[ChatMessage]) -> L
# Abstract Model Backend
# =============================================================================
class ModelBackend(ABC):
"""Abstract base class for model backends."""
@abstractmethod
def load_model(self, model_name: str, **kwargs) -> None:
"""Load the model."""
pass
@abstractmethod
def generate(self, prompt: str, max_tokens: Optional[int] = None,
temperature: float = 0.7, top_p: float = 1.0,
stop: Optional[List[str]] = None) -> str:
"""Generate text non-streaming."""
pass
@abstractmethod
def generate_stream(self, prompt: str, max_tokens: Optional[int] = None,
temperature: float = 0.7, top_p: float = 1.0,
stop: Optional[List[str]] = None) -> AsyncGenerator[str, None]:
"""Generate text in streaming fashion."""
pass
@abstractmethod
def format_messages(self, messages: List[ChatMessage]) -> str:
"""Format messages into a prompt string."""
pass
@abstractmethod
def get_model_name(self) -> str:
"""Return the loaded model name."""
pass
@abstractmethod
def cleanup(self) -> None:
"""Cleanup resources."""
pass
# =============================================================================
# NVIDIA/HuggingFace Backend
# =============================================================================
......@@ -2486,216 +2215,6 @@ class VulkanBackend(ModelBackend):
# Model Manager
# =============================================================================
class ModelManager:
"""Manages the loaded model and tokenizer."""
def __init__(self):
self.backend: Optional[ModelBackend] = None
self.backend_type: Optional[str] = None
self.tool_parser = ModelParserAdapter()
def _aggressive_vram_cleanup(self, model_manager):
"""
Aggressively cleanup VRAM when switching between different model types.
This is more thorough than a simple cleanup() call.
"""
import gc
import time
try:
import torch
# First, try to move model to CPU if it has a model attribute
if hasattr(model_manager, 'model') and model_manager.model is not None:
model = model_manager.model
# If it's a diffusers pipeline, try to move to CPU first
if hasattr(model, 'to'):
try:
model.to('cpu')
except:
pass
# Delete the model
del model
# Also handle backend directly if it's different
if hasattr(model_manager, 'backend') and model_manager.backend is not None:
backend = model_manager.backend
if hasattr(backend, 'model') and backend.model is not None:
model = backend.model
if hasattr(model, 'to'):
try:
model.to('cpu')
except:
pass
del model
if hasattr(backend, 'pipeline') and backend.pipeline is not None:
del backend.pipeline
if hasattr(backend, 'vae') and backend.vae is not None:
del backend.vae
if hasattr(backend, 'text_encoder') and backend.text_encoder is not None:
del backend.text_encoder
if hasattr(backend, 'tokenizer') and backend.tokenizer is not None:
del backend.tokenizer
# Force multiple rounds of garbage collection
for _ in range(3):
gc.collect()
# Clear PyTorch cache
if torch.cuda.is_available():
torch.cuda.synchronize()
torch.cuda.empty_cache()
# Add delay to allow Vulkan to release memory
time.sleep(2)
except Exception as e:
print(f"Warning during aggressive VRAM cleanup: {e}")
finally:
# Try to cleanup the model manager itself
try:
if hasattr(model_manager, 'cleanup'):
model_manager.cleanup()
except:
pass
def load_model(self, model_name: str, backend_type: str = "auto", **kwargs):
"""
Load the model with the specified backend.
Args:
model_name: Model name or path
backend_type: 'nvidia', 'vulkan', 'cuda', or 'auto' to detect
**kwargs: Additional arguments for the specific backend
"""
available = detect_available_backends()
# Check if model is a GGUF file
is_gguf = model_name.endswith('.gguf') or 'gguf' in model_name.lower()
# Determine backend
if backend_type == "auto":
if available.get('nvidia'):
backend_type = "nvidia"
print("Auto-detected NVIDIA backend")
elif available.get('vulkan'):
backend_type = "vulkan"
print("Auto-detected Vulkan backend")
else:
print("No GPU backend detected. For NVIDIA, install PyTorch with CUDA.")
print("For Vulkan, install llama-cpp-python with Vulkan support.")
raise RuntimeError("No suitable backend found")
# If GGUF file and backend is nvidia/cuda, use llama-cpp-python with CUDA backend
original_backend = None
if is_gguf and backend_type in ("nvidia", "cuda"):
original_backend = backend_type
print(f"GGUF model detected, using llama-cpp-python (original backend: {original_backend})")
backend_type = "vulkan" # Use llama-cpp-python for GGUF
self.backend_type = backend_type
# Create appropriate backend
if backend_type == "nvidia":
if not available.get('nvidia'):
raise RuntimeError("NVIDIA backend requested but PyTorch/CUDA not available")
self.backend = NvidiaBackend()
elif backend_type == "vulkan":
if not available.get('vulkan'):
raise RuntimeError("Vulkan backend requested but llama-cpp-python not available")
self.backend = VulkanBackend(original_backend=original_backend)
else:
raise ValueError(f"Unknown backend: {backend_type}")
# Load the model
self.backend.load_model(model_name, **kwargs)
self.tool_parser = ModelParserAdapter(model_name=model_name)
def format_messages(self, messages: List[ChatMessage]) -> str:
"""Format messages into a prompt string."""
if self.backend is None:
raise RuntimeError("No model loaded")
return self.backend.format_messages(messages)
def generate(self, prompt: str, max_tokens: Optional[int] = None,
temperature: float = 0.7, top_p: float = 1.0,
stop: Optional[List[str]] = None) -> str:
"""Generate text non-streaming."""
if self.backend is None:
raise RuntimeError("No model loaded")
return self.backend.generate(prompt, max_tokens, temperature, top_p, stop)
def generate_chat(self, messages: List[Dict], max_tokens: Optional[int] = None,
temperature: float = 0.7, top_p: float = 1.0,
stop: Optional[List[str]] = None, tools: Optional[List] = None,
response_format: Optional[Dict] = None) -> str:
"""Generate chat completion non-streaming."""
if self.backend is None:
raise RuntimeError("No model loaded")
# Use generate_chat if available (Vulkan backend), otherwise format and use generate
if hasattr(self.backend, 'generate_chat'):
return self.backend.generate_chat(messages, max_tokens, temperature, top_p, stop, tools, response_format)
else:
# Fallback for NVIDIA backend
prompt = self.format_messages([ChatMessage(**m) for m in messages])
return self.backend.generate(prompt, max_tokens, temperature, top_p, stop)
async def generate_stream(self, prompt: str, max_tokens: Optional[int] = None,
temperature: float = 0.7, top_p: float = 1.0,
stop: Optional[List[str]] = None) -> AsyncGenerator[str, None]:
"""Generate text in streaming fashion."""
if self.backend is None:
raise RuntimeError("No model loaded")
async for chunk in self.backend.generate_stream(prompt, max_tokens, temperature, top_p, stop):
yield chunk
async def generate_chat_stream(self, messages: List[Dict], max_tokens: Optional[int] = None,
temperature: float = 0.7, top_p: float = 1.0,
stop: Optional[List[str]] = None, tools: Optional[List] = None,
response_format: Optional[Dict] = None) -> AsyncGenerator[str, None]:
"""Generate chat completion streaming."""
if self.backend is None:
raise RuntimeError("No model loaded")
# Use generate_chat_stream if available (Vulkan backend), otherwise format and use generate_stream
if hasattr(self.backend, 'generate_chat_stream'):
async for chunk in self.backend.generate_chat_stream(messages, max_tokens, temperature, top_p, stop, tools, response_format):
yield chunk
else:
# Fallback for NVIDIA backend
prompt = self.format_messages([ChatMessage(**m) for m in messages])
async for chunk in self.backend.generate_stream(prompt, max_tokens, temperature, top_p, stop):
yield chunk
@property
def model_name(self) -> str:
if self.backend is None:
return "unknown"
return self.backend.get_model_name()
@property
def model(self):
if self.backend is None:
return None
return self.backend
@property
def tokenizer(self):
# Only NVIDIA backend has a tokenizer
if isinstance(self.backend, NvidiaBackend):
return self.backend.tokenizer
return None
def cleanup(self):
if self.backend is not None:
self.backend.cleanup()
self.backend = None
# =============================================================================
# Whisper Server Manager - manages whisper-server subprocess
# =============================================================================
......@@ -2705,803 +2224,10 @@ import signal
import requests
import time
import threading
class WhisperServerManager:
"""Manages whisper-server subprocess for audio transcription with model swapping support."""
def __init__(self, server_path: str = None, port: int = 8744):
self.server_path = server_path
self.port = port
self.process = None
self.current_model = None
self.base_url = f"http://127.0.0.1:{port}"
self.lock = threading.Lock()
self._health_check_thread = None
self._running = False
# Check if port is available
if not self._is_port_available(port):
# Try to find an available port
for new_port in range(port + 1, port + 100):
if self._is_port_available(new_port):
self.port = new_port
self.base_url = f"http://127.0.0.1:{new_port}"
print(f"Port {port} in use, using port {new_port} instead")
break
def _is_port_available(self, port: int) -> bool:
"""Check if a port is available."""
import socket
try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(('127.0.0.1', port))
return True
except OSError:
return False
def is_running(self) -> bool:
"""Check if whisper-server is running."""
if self.process is None:
return False
return self.process.poll() is None
def start(self, model_path: str = None, gpu_device: int = 0) -> str:
"""Start whisper-server with the specified model. Returns actual model path or empty string on failure."""
with self.lock:
# Stop existing server if running
if self.is_running():
self.stop()
if not self.server_path:
print("Error: whisper-server path not set")
return ""
# Handle URL models - download if needed
actual_model_path = model_path
if model_path and (model_path.startswith('http://') or model_path.startswith('https://')):
# Check cache first
cached_path = get_cached_model_path(model_path)
if cached_path:
actual_model_path = cached_path
print(f"Using cached model: {actual_model_path}")
else:
# Download the model
cache_dir = get_model_cache_dir()
print(f"Downloading model: {model_path}")
actual_model_path = download_model(model_path, cache_dir)
print(f"Downloaded model to: {actual_model_path}")
# Build command
cmd = [self.server_path]
if actual_model_path:
cmd.extend(["-m", actual_model_path])
# Add GPU device
cmd.extend(["-dev", str(gpu_device)])
# Add --convert flag to convert audio to 16kHz mono on the server side
cmd.append("--convert")
# Add host and port
cmd.extend(["--host", "127.0.0.1"])
cmd.extend(["--port", str(self.port)])
print(f"Starting whisper-server: {' '.join(cmd)}")
print(f"DEBUG: Full whisper-server command: {' '.join(cmd)}")
try:
self.process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
preexec_fn=lambda: signal.signal(signal.SIGTERM, signal.SIG_DFL)
)
self.current_model = actual_model_path
# Wait for server to be ready
if self._wait_for_server(30):
print(f"whisper-server started on {self.base_url}")
self._running = True
return actual_model_path
else:
print("Error: whisper-server failed to start")
self.stop()
return ""
except Exception as e:
print(f"Error starting whisper-server: {e}")
return ""
def stop(self):
"""Stop whisper-server."""
with self.lock:
self._running = False
if self.process:
try:
self.process.terminate()
try:
self.process.wait(timeout=5)
except subprocess.TimeoutExpired:
self.process.kill()
self.process.wait()
except Exception as e:
print(f"Error stopping whisper-server: {e}")
self.process = None
self.current_model = None
def restart(self, model_path: str = None, gpu_device: int = 0) -> bool:
"""Restart whisper-server with a new model."""
print(f"Restarting whisper-server with model: {model_path}")
return self.start(model_path, gpu_device)
def transcribe(self, audio_data: bytes, language: str = None, prompt: str = None) -> dict:
"""Send transcription request to whisper-server."""
if not self.is_running():
return {"error": "whisper-server not running"}
try:
files = {"file": ("audio.wav", audio_data, "audio/wav")}
data = {}
if language:
data["language"] = language
if prompt:
data["prompt"] = prompt
print(f"DEBUG: Sending POST to {self.base_url}/inference with data={data}, file_size={len(audio_data)}")
response = requests.post(
f"{self.base_url}/inference",
files=files,
data=data,
timeout=300
)
print(f"DEBUG: whisper-server response status={response.status_code}, body={response.text[:500] if response.text else 'empty'}")
if response.status_code == 200:
return response.json()
else:
return {"error": f"Server error: {response.status_code}", "detail": response.text}
except Exception as e:
print(f"DEBUG: whisper-server exception: {e}")
return {"error": str(e)}
def _wait_for_server(self, timeout: int = 30) -> bool:
"""Wait for whisper-server to be ready."""
start_time = time.time()
while time.time() - start_time < timeout:
try:
response = requests.get(f"{self.base_url}/health", timeout=2)
if response.status_code == 200:
return True
except:
pass
time.sleep(0.5)
return False
def get_status(self) -> dict:
"""Get whisper-server status."""
return {
"running": self.is_running(),
"model": self.current_model,
"url": self.base_url
}
# =============================================================================
# Multi-Model Manager (supports audio transcription and image generation)
# =============================================================================
class MultiModelManager:
"""
Manages multiple models: main text model, audio transcription, and image generation.
Supports dynamic switching based on request model name.
Modes:
- default: Load models on-demand (swap models in VRAM when request changes)
- loadall: Pre-load all models in VRAM at startup
- loadswap: Keep all models in memory (CPU RAM), swap active model to VRAM
"""
def __init__(self):
self.models: Dict[str, ModelManager] = {}
self.default_model: Optional[str] = None
self.audio_models: List[str] = [] # List of audio model names
self.tts_model: Optional[str] = None
self.image_models: List[str] = [] # List of image model names
self.vision_models: List[str] = [] # List of vision (image/video to text) model names
self.tool_parser = ModelParserAdapter()
self.current_model_key: Optional[str] = None
# Configuration for each model type
self.config: Dict[str, Dict] = {}
# Load mode settings
self.load_mode: str = "ondemand" # "ondemand", "loadall", "loadswap"
self.active_in_vram: Optional[str] = None # Which model is currently in VRAM
# Model aliases: alias -> actual model name mapping
self.model_aliases: Dict[str, str] = {}
# Whisper server manager
self.whisper_server: Optional[WhisperServerManager] = None
# Track backend type for each model (needed for on-demand loading)
self.model_backend_types: Dict[str, str] = {}
def _aggressive_vram_cleanup(self, model_manager):
"""
Aggressively cleanup VRAM when switching between different model types.
This is more thorough than a simple cleanup() call.
"""
import gc
import time
try:
import torch
# First, try to move model to CPU if it has a model attribute
if hasattr(model_manager, 'model') and model_manager.model is not None:
model = model_manager.model
# If it's a diffusers pipeline, try to move to CPU first
if hasattr(model, 'to'):
try:
model.to('cpu')
except:
pass
# Delete the model
del model
# Also handle backend directly if it's different
if hasattr(model_manager, 'backend') and model_manager.backend is not None:
backend = model_manager.backend
if hasattr(backend, 'model') and backend.model is not None:
model = backend.model
if hasattr(model, 'to'):
try:
model.to('cpu')
except:
pass
del model
if hasattr(backend, 'pipeline') and backend.pipeline is not None:
del backend.pipeline
if hasattr(backend, 'vae') and backend.vae is not None:
del backend.vae
if hasattr(backend, 'text_encoder') and backend.text_encoder is not None:
del backend.text_encoder
if hasattr(backend, 'tokenizer') and backend.tokenizer is not None:
del backend.tokenizer
# Force multiple rounds of garbage collection
for _ in range(3):
gc.collect()
# Clear PyTorch cache
if torch.cuda.is_available():
torch.cuda.synchronize()
torch.cuda.empty_cache()
# Add delay to allow Vulkan to release memory
time.sleep(2)
except Exception as e:
print(f"Warning during aggressive VRAM cleanup: {e}")
finally:
# Try to cleanup the model manager itself
try:
if hasattr(model_manager, 'cleanup'):
model_manager.cleanup()
except:
pass
# Load mode settings
self.load_mode: str = "ondemand" # "ondemand", "loadall", "loadswap"
self.active_in_vram: Optional[str] = None # Which model is currently in VRAM
# Model aliases: alias -> actual model name mapping
self.model_aliases: Dict[str, str] = {}
# Whisper server manager
self.whisper_server: Optional[WhisperServerManager] = None
# Track backend type for each model (needed for on-demand loading)
self.model_backend_types: Dict[str, str] = {}
@property
def audio_model(self) -> Optional[str]:
"""Get the first/default audio model."""
return self.audio_models[0] if self.audio_models else None
@property
def image_model(self) -> Optional[str]:
"""Get the first/default image model."""
return self.image_models[0] if self.image_models else None
@property
def vision_model(self) -> Optional[str]:
"""Get the first/default vision model."""
return self.vision_models[0] if self.vision_models else None
def set_load_mode(self, mode: str):
"""Set the load mode: 'ondemand', 'loadall', or 'loadswap'."""
self.load_mode = mode
def set_default_model(self, model_name: str, config: Dict = None, backend_type: str = "auto"):
"""Set the default/main text model."""
self.default_model = model_name
self.config[model_name] = config or {}
self.model_backend_types[model_name] = backend_type
def set_audio_model(self, model_name: str, config: Dict = None):
"""Add an audio transcription model."""
if model_name not in self.audio_models:
self.audio_models.append(model_name)
self.config[f"audio:{model_name}"] = config or {}
def set_tts_model(self, model_name: str, config: Dict = None):
"""Set the text-to-speech model."""
self.tts_model = model_name
self.config[f"tts:{model_name}"] = config or {}
def set_image_model(self, model_name: str, config: Dict = None):
"""Add an image generation model."""
if model_name not in self.image_models:
self.image_models.append(model_name)
self.config[f"image:{model_name}"] = config or {}
def set_vision_model(self, model_name: str, config: Dict = None):
"""Add a vision (image/video to text) model."""
if model_name not in self.vision_models:
self.vision_models.append(model_name)
self.config[f"vision:{model_name}"] = config or {}
def set_model_alias(self, alias: str, model_name: str):
"""Register an alias for a model."""
self.model_aliases[alias] = model_name
def get_model_for_request(self, requested_model: str) -> Optional[ModelManager]:
"""
Get the appropriate model manager for a request based on model name.
Model name conventions:
- "default", empty, or matches default model -> use main model
- "audio" -> use first/default audio model
- "audio:modelname" -> use specific audio model
- "image" -> use first/default image model
- "vision:modelname" or "image:modelname" -> use specific image model
- "tts" -> use TTS model
- "tts:modelname" -> use specific TTS model
- Custom aliases -> resolve to actual model name
- Otherwise match by model ID in multi_model_manager.models
In ondemand mode with multiple text models:
- If requested model is different from currently loaded model,
unload current and load new model on-demand (respecting --backend)
"""
# Import global_args inside function to ensure it's available
global global_args
# Resolve custom aliases first
if requested_model in self.model_aliases:
requested_model = self.model_aliases[requested_model]
# Handle empty or "default" model names
if not requested_model or requested_model == "default":
if self.default_model and self.default_model in self.models:
self.current_model_key = self.default_model
return self.models[self.default_model]
# Model not loaded - check if it's in config (registered but unloaded)
if self.default_model and self.default_model in self.config:
# Need to reload the default model - cleanup image models first
for key in list(self.models.keys()):
if key.startswith("image:"):
model_to_cleanup = self.models.get(key)
if model_to_cleanup is not None:
print(f"Unloading image model '{key}' from VRAM to reload text model")
self._aggressive_vram_cleanup(model_to_cleanup)
del self.models[key]
# Add delay to allow VRAM to be freed
import time
time.sleep(2)
# Now try to reload the default model
try:
from llama_cpp import Llama
backend = self.config[self.default_model].get('backend_type', 'auto')
model_path = self.default_model
# Check if model_path is a URL and try to get cached path
if model_path.startswith('http://') or model_path.startswith('https://'):
cached_path = get_cached_model_path(model_path)
if cached_path:
model_path = cached_path
print(f"Using cached model path: {model_path}")
else:
print(f"Warning: Model URL not cached, cannot reload: {model_path}")
return None
load_kwargs = self.config[self.default_model].copy()
load_kwargs.pop('backend_type', None)
print(f"Reloading default model: {model_path}")
llm = Llama(model_path=model_path, **load_kwargs)
self.models[self.default_model] = ModelManager(llm, backend=backend)
self.current_model_key = self.default_model
return self.models[self.default_model]
except Exception as e:
print(f"Error reloading default model: {e}")
return None
# Handle "audio" alias - use first/default audio model
if requested_model == "audio":
if self.audio_models:
first_audio = self.audio_models[0]
key = f"audio:{first_audio}"
if key in self.models:
self.current_model_key = key
return self.models[key]
# Try to load on demand
return None
return None
# Handle "image" alias - use first/default image model
if requested_model == "image":
if self.image_models:
first_image = self.image_models[0]
key = f"image:{first_image}"
if key in self.models:
self.current_model_key = key
return self.models[key]
# Try to load on demand
return None
return None
# Handle "tts" alias
if requested_model == "tts":
if self.tts_model:
key = f"tts:{self.tts_model}"
if key in self.models:
self.current_model_key = key
return self.models[key]
return None
return None
# Check for specialized models with prefix
if requested_model.startswith("audio:"):
audio_name = requested_model[6:] # Remove "audio:" prefix
key = f"audio:{audio_name}"
if key in self.models:
self.current_model_key = key
return self.models[key]
elif audio_name in self.audio_models:
# Try loading audio model on demand
return None
return None
if requested_model.startswith("tts:"):
tts_name = requested_model[4:] # Remove "tts:" prefix
key = f"tts:{tts_name}"
if key in self.models:
self.current_model_key = key
return self.models[key]
elif self.tts_model and tts_name == self.tts_model:
return None
return None
# Handle both "vision:" and "image:" prefixes
if requested_model.startswith("vision:") or requested_model.startswith("image:"):
# Extract the model name (remove either prefix)
if requested_model.startswith("vision:"):
image_name = requested_model[7:] # Remove "vision:" prefix
else:
image_name = requested_model[6:] # Remove "image:" prefix
key = f"image:{image_name}"
if key in self.models:
self.current_model_key = key
return self.models[key]
elif image_name in self.image_models:
# Try loading image model on demand
return None
return None
# Check if it's the default model
if self.default_model and (requested_model == self.default_model or
requested_model.endswith(self.default_model.split("/")[-1])):
self.current_model_key = self.default_model
return self.models.get(self.default_model)
# Check if any loaded model matches
for key, model in self.models.items():
if requested_model in key or key.endswith(requested_model.split("/")[-1]):
self.current_model_key = key
return model
# === ON-DEMAND MODEL SWITCHING FOR TEXT MODELS ===
# If we're in ondemand mode and the requested model is in config but not loaded,
# we should try to load it on-demand (swap from current model)
# Only for text models (not audio/image/tts which have their own handling)
# First, cleanup any image models to free VRAM for text model
for key in list(self.models.keys()):
if key.startswith("image:"):
model_to_cleanup = self.models.get(key)
if model_to_cleanup is not None:
print(f"Unloading image model '{key}' from VRAM to make room for text model")
self._aggressive_vram_cleanup(model_to_cleanup)
del self.models[key]
# Force garbage collection and clear GPU cache
import gc
gc.collect()
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except:
pass
# Add a longer delay to allow VRAM to be freed (Vulkan needs more time)
import time
time.sleep(2)
# Check if requested model is already loaded - if so, reuse it
if requested_model in self.models:
self.current_model_key = requested_model
return self.models[requested_model]
# Check if requested model is in our config (means it was registered but not loaded)
if self.load_mode == "ondemand" and requested_model in self.config:
# This is a text model that's registered but not loaded
# We need to swap: unload current model and load this one
# Always cleanup any loaded model (unless it's the same model we're about to load)
for key in list(self.models.keys()):
if key != requested_model:
model_to_cleanup = self.models.get(key)
if model_to_cleanup is not None:
print(f"Unloading '{key}' from VRAM to load '{requested_model}'")
try:
if hasattr(model_to_cleanup, 'cleanup') and callable(getattr(model_to_cleanup, 'cleanup')):
model_to_cleanup.cleanup()
elif hasattr(model_to_cleanup, 'model') and model_to_cleanup.model is not None:
# Handle ModelManager objects
if hasattr(model_to_cleanup.model, 'cleanup'):
model_to_cleanup.model.cleanup()
except Exception as e:
print(f"Warning during cleanup of '{key}': {e}")
del self.models[key]
# Load the new model on-demand
print(f"ON-DEMAND SWAP: Loading model '{requested_model}' into VRAM")
# Get the backend type for this model
backend_type = getattr(self, 'model_backend_types', {}).get(requested_model, "auto")
# Get config for this model
model_config = self.config.get(requested_model, {})
effective_backend = backend_type
if effective_backend == "auto" and global_args:
effective_backend = getattr(global_args, 'backend', 'auto')
try:
# Create new model manager and load the model
new_manager = ModelManager()
new_manager.load_model(
model_name=requested_model,
backend_type=effective_backend,
**model_config
)
self.models[requested_model] = new_manager
self.current_model_key = requested_model
self.active_in_vram = requested_model
print(f"ON-DEMAND SWAP: Successfully loaded model '{requested_model}' with backend '{effective_backend}'")
return new_manager
except Exception as e:
print(f"ON-DEMAND SWAP: Failed to load model '{requested_model}': {e}")
# Try to restore the previous model if we had one
return None
# Also check if the model matches by short name (e.g., "Phi-3" matches "microsoft/Phi-3-mini-4k-instruct")
if self.load_mode == "ondemand":
# First, cleanup any image models to free VRAM for text model
for key in list(self.models.keys()):
if key.startswith("image:"):
model_to_cleanup = self.models.get(key)
if model_to_cleanup is not None:
print(f"Unloading image model '{key}' from VRAM to make room for text model")
try:
if hasattr(model_to_cleanup, 'cleanup') and callable(getattr(model_to_cleanup, 'cleanup')):
model_to_cleanup.cleanup()
elif hasattr(model_to_cleanup, 'model') and model_to_cleanup.model is not None:
if hasattr(model_to_cleanup.model, 'cleanup'):
model_to_cleanup.model.cleanup()
except Exception as e:
print(f"Warning during cleanup of '{key}': {e}")
del self.models[key]
for model_name in self.config.keys():
# Only check text models (not audio:, image:, tts: prefixes)
if ":" not in model_name:
short_name = model_name.split("/")[-1] if "/" in model_name else model_name
if requested_model.lower() in short_name.lower() or short_name.lower() in requested_model.lower():
# Found a matching model in config, try to load it
if model_name not in self.models:
# Always cleanup any loaded model (unless it's the same model we're about to load)
for key in list(self.models.keys()):
if key != model_name:
model_to_cleanup = self.models.get(key)
if model_to_cleanup is not None:
print(f"Unloading '{key}' from VRAM to load '{model_name}'")
try:
if hasattr(model_to_cleanup, 'cleanup') and callable(getattr(model_to_cleanup, 'cleanup')):
model_to_cleanup.cleanup()
elif hasattr(model_to_cleanup, 'model') and model_to_cleanup.model is not None:
if hasattr(model_to_cleanup.model, 'cleanup'):
model_to_cleanup.model.cleanup()
except Exception as e:
print(f"Warning during cleanup of '{key}': {e}")
del self.models[key]
# Load the new model on-demand
print(f"ON-DEMAND SWAP: Loading model '{model_name}' into VRAM")
# Get the backend type for this model
backend_type = getattr(self, 'model_backend_types', {}).get(model_name, "auto")
# Get config for this model
model_config = self.config.get(model_name, {})
effective_backend = backend_type
if effective_backend == "auto" and global_args:
effective_backend = getattr(global_args, 'backend', 'auto')
try:
new_manager = ModelManager()
new_manager.load_model(
model_name=model_name,
backend_type=effective_backend,
**model_config
)
self.models[model_name] = new_manager
self.current_model_key = model_name
self.active_in_vram = model_name
print(f"ON-DEMAND SWAP: Successfully loaded model '{model_name}' with backend '{effective_backend}'")
return new_manager
except Exception as e:
print(f"ON-DEMAND SWAP: Failed to load model '{model_name}': {e}")
return None
return None
def add_model(self, key: str, manager: ModelManager):
"""Add a model manager for a specific key."""
self.models[key] = manager
def get_model(self, key: str) -> Optional[ModelManager]:
"""Get a model manager by key."""
return self.models.get(key)
def get_current_model(self) -> Optional[ModelManager]:
"""Get the currently active model."""
if self.current_model_key:
return self.models.get(self.current_model_key)
if self.default_model:
return self.models.get(self.default_model)
return None
def list_models(self) -> List[ModelInfo]:
"""List all available models."""
models = []
# Add default model(s)
if self.default_model:
model_id = self.default_model
# Skip URLs - they are download sources, not model identifiers
if not (model_id.startswith("http://") or model_id.startswith("https://")):
# Also add short name
short_name = self.default_model.split("/")[-1] if "/" in self.default_model else self.default_model
if short_name != self.default_model:
models.append(ModelInfo(id=short_name))
models.append(ModelInfo(id=model_id))
models.append(ModelInfo(id="default"))
# Add aliases for first/default models
if self.audio_models:
models.append(ModelInfo(id="audio")) # Alias for first audio model
# Add all audio models
for audio_id in self.audio_models:
models.append(ModelInfo(id=f"audio:{audio_id}"))
# Add TTS models
if self.tts_model:
models.append(ModelInfo(id="tts")) # Alias for TTS
tts_id = f"tts:{self.tts_model}"
models.append(ModelInfo(id=tts_id))
# Add vision/image models
if self.image_models:
models.append(ModelInfo(id="image")) # Alias for first image model
# Add all image models - convert URLs to cached paths for display
for image_id in self.image_models:
# Check if image_id is a URL and try to get cached path
display_id = image_id
if image_id.startswith('http://') or image_id.startswith('https://'):
cached_path = get_cached_model_path(image_id)
if cached_path:
# Use the filename from cached path for display
display_id = os.path.basename(cached_path)
models.append(ModelInfo(id=f"image:{display_id}"))
# Add loaded models that aren't in the above categories
for key in self.models:
# Skip if already added
if key == self.default_model or key.startswith("audio:") or key.startswith("image:") or key.startswith("tts:"):
continue
# Skip short names (already added)
if self.default_model and key == self.default_model.split("/")[-1]:
continue
# Skip URLs - they are download sources, not model identifiers
if key.startswith("http://") or key.startswith("https://"):
continue
models.append(ModelInfo(id=key))
# Add custom model aliases
for alias in self.model_aliases:
models.append(ModelInfo(id=alias))
return models if models else [ModelInfo(id="default")]
def swap_model_to_vram(self, model_key: str) -> bool:
"""
Swap a model from CPU RAM to GPU VRAM.
Returns True if successful, False otherwise.
Only applies in loadswap mode.
"""
if self.load_mode != "loadswap":
return True # No swapping needed in other modes
if self.active_in_vram == model_key:
return True # Already in VRAM
# If another model is in VRAM, swap it to CPU first
if self.active_in_vram and self.active_in_vram in self.models:
self._swap_model_to_cpu(self.active_in_vram)
# Swap the requested model to VRAM
success = self._swap_model_to_vram(model_key)
if success:
self.active_in_vram = model_key
return success
def _swap_model_to_vram(self, model_key: str) -> bool:
"""Internal method to swap a specific model to VRAM."""
# This would need backend-specific implementation
# For now, we assume the model is already in memory
print(f"SWAP: Moving model {model_key} to VRAM")
return True
def _swap_model_to_cpu(self, model_key: str) -> bool:
"""Internal method to swap a specific model to CPU RAM."""
# This would need backend-specific implementation
# For now, we just track that it's in CPU
print(f"SWAP: Moving model {model_key} to CPU RAM")
return True
def get_active_model_key(self) -> Optional[str]:
"""Get the currently active model key."""
return self.current_model_key or self.default_model
def cleanup(self):
"""Cleanup all models."""
for key, model in self.models.items():
try:
if hasattr(model, 'cleanup') and callable(getattr(model, 'cleanup')):
model.cleanup()
elif hasattr(model, 'model') and model.model is not None:
# Handle ModelManager objects
if hasattr(model.model, 'cleanup'):
model.model.cleanup()
elif hasattr(model.model, 'model'):
# Nested model (e.g., StableDiffusion wrapper)
if model.model.model is not None:
del model.model.model
# Remove from dict
del self.models[key]
except Exception as e:
print(f"Warning: Failed to cleanup model '{key}': {e}")
self.models.clear()
# Global multi-model manager
multi_model_manager = MultiModelManager()
# Global model manager (for backward compatibility)
......@@ -3511,227 +2237,11 @@ model_manager = ModelManager()
global_args = None
def check_hf_chat_template(model_type: str = "text", model_name: str = None) -> tuple:
"""
Check if HuggingFace chat template should be used for the model.
Returns a tuple (should_use, template_name) where template_name is the template to use or None for auto-detect.
Args:
model_type: The model type ('text', 'image', etc.)
model_name: The specific model name (optional)
Returns:
Tuple of (should_use: bool, template_name: str or None)
template_name is None means auto-detect from tokenizer
Examples:
# Auto-detect and apply to all models
--hf-chat-template auto
# Apply to all text models with auto-detect
--hf-chat-template text
# Apply to specific model with auto-detect
--hf-chat-template text:llama-3.1
# Apply to specific model with specific template
--hf-chat-template "llama-3.1:llama3"
--hf-chat-template "phi-3:chatml"
"""
hf_chat_template = getattr(global_args, 'hf_chat_template', []) or []
# If empty list, HF chat template is not enabled
if not hf_chat_template:
return (False, None)
for spec in hf_chat_template:
# Handle auto-detect - try to load HF tokenizer and auto-detect template
if spec == 'auto' or spec == '':
# Applies to all models when using 'auto'
return (True, None)
# Check if this spec has a template specified after the model name
# Format: "model_name:template_name" or "type:model_name:template_name"
parts = spec.split(':')
if len(parts) == 1:
# Just a type or single value
spec_val = parts[0]
if spec_val == model_type or spec_val == '*':
return (True, None)
# Check if it matches the model name directly (when model_type is part of the name)
if model_name and (spec_val in model_name or model_name in spec_val):
return (True, None)
elif len(parts) == 2:
# Format: "type:model_name" or "model_name:template"
spec_type = parts[0]
spec_model = parts[1]
# Check if it's "text" or "image" type
if spec_type in ('text', 'image', '*'):
if spec_type == model_type or spec_type == '*':
# Check if model name matches
if spec_model == model_name or spec_model == '*':
return (True, None)
else:
# It's "model_name:template" format
if model_name and (spec_model in model_name or model_name in spec_model):
return (True, spec_type) # spec_type is actually the template!
elif len(parts) == 3:
# Format: "type:model_name:template"
spec_type = parts[0]
spec_model = parts[1]
spec_template = parts[2]
if spec_type == model_type or spec_type == '*':
if spec_model == model_name or spec_model == '*':
return (True, spec_template)
return (False, None)
# Global system prompt (set via --system-prompt flag)
# None = don't inject, True = use default, string = use custom text
global_system_prompt = None
def get_resolved_model_name(requested_model: str, current_manager = None) -> str:
"""
Get the actual model name to return in the response.
- If the requested model is an alias (like "default"), resolve to actual model name
- For URL models, extract just the filename
- Add "coderai/" as the provider prefix
"""
import os
# First try to get from model manager if available
if current_manager is not None:
actual_model = current_manager.model_name
if actual_model and actual_model != "unknown":
model_name = actual_model
else:
model_name = requested_model
else:
model_name = requested_model
# Handle URL models - extract just the filename
if model_name.startswith('http://') or model_name.startswith('https://'):
# Extract filename from URL
url_path = model_name.split('?')[0] # Remove query params
filename = os.path.basename(url_path)
if filename:
model_name = filename
# Handle local file paths - extract just the filename
elif os.path.isfile(model_name):
model_name = os.path.basename(model_name)
# Add coderai/ prefix for the response
if not model_name.startswith("coderai/"):
model_name = f"coderai/{model_name}"
return model_name
def get_model_family(model_name: str) -> str:
"""Detect model family from model name."""
model_lower = model_name.lower()
if 'qwen' in model_lower:
return 'qwen'
elif 'deepseek' in model_lower:
return 'deepseek'
elif 'llama-3' in model_lower or 'llama3' in model_lower or 'meta-llama-3' in model_lower:
return 'llama3'
elif 'llama' in model_lower or 'meta-llama' in model_lower:
return 'llama'
elif 'mistral' in model_lower or 'mixtral' in model_lower:
return 'mistral'
elif 'gemma' in model_lower:
return 'gemma'
elif 'yi' in model_lower:
return 'yi'
elif 'hermes' in model_lower:
return 'hermes'
else:
return 'unknown'
def get_reasoning_stop_tokens(model_family: str) -> tuple:
"""Get stop tokens for reasoning mode based on model family.
Returns tuple of (start_token, end_token, additional_stops)
"""
if model_family == 'qwen':
# Qwen uses <|im_start|> format with </tool_call> for thinking
return (
"<|im_start|>assistant\n",
"<|im_end|>",
["<|im_end|>", "<|endoftext|>"]
)
elif model_family == 'deepseek':
return (
"<｜Assistant｜>",
"<｜endofassistant｜>",
["<｜endofassistant｜>", "<｜User｜>", "<｜endoftext｜>"]
)
elif model_family == 'llama3':
return (
"<|start_header_id|>assistant<|end_header_id|>\n\n<thought>\n",
"</thought>",
["</thought>", "<|eot_id|>", "<|end_of_text|>"]
)
elif model_family == 'llama':
return (
"<|start_header_id|>assistant<|end_header_id|>\n\n",
"<|eot_id|>",
["<|eot_id|>", "<|end_of_text|>"]
)
elif model_family == 'mistral':
return (
"[/INST] <thought>\n",
"</thought>",
["</thought>", "</INST>", "[INST]"]
)
elif model_family == 'gemma':
return (
"<start_of_turn>model\n<thought>\n",
"</thought>",
["</thought>", "<end_of_turn>", "<start_of_turn>"]
)
elif model_family == 'yi' or model_family == 'hermes':
return (
"<|im_start|>assistant\n",
"<|im_end|>",
["<|im_end|>", "<|endoftext|>"]
)
else:
# Default fallback - try common tokens
return (
"<|im_start|>assistant\n",
"<|im_end|>",
["<|im_end|>", "<|endoftext|>"]
)
def get_reasoning_system_prompt(model_family: str) -> str:
"""Get system prompt injection for forcing reasoning on non-native models."""
if model_family == 'qwen':
return "You must reason step-by-step inside <thought> tags before every response."
elif model_family == 'deepseek':
return "You must reason step-by-step inside <0x00></think> tags before every response."
elif model_family in ('llama3', 'llama'):
return "You must reason step-by-step inside <thought> tags before every response."
elif model_family == 'mistral':
return "You must reason step-by-step inside <thought> tags before every response."
elif model_family == 'gemma':
return "You must reason step-by-step inside <thought> tags before every response. Use <start_of_turn>model for your response."
elif model_family in ('yi', 'hermes'):
return "You must reason step-by-step inside <|im_start|>assistant tags before every response."
else:
return "You must reason step-by-step before every response."
# Global debug flag
global_debug = False
global_file_path = None
......@@ -3740,61 +2250,6 @@ global_file_path = None
# Queue Manager for Model Loading Notifications
# =============================================================================
class QueueManager:
"""
Manages request queue for model loading notifications.
When clients are waiting for a model to load, sends them progress updates.
"""
def __init__(self):
self.waiting_requests: Dict[str, float] = {} # request_id -> start_time
self.current_request_id: Optional[str] = None
self.model_loading: bool = False
self.model_name: Optional[str] = None
self.lock = asyncio.Lock()
async def add_waiting(self, request_id: str) -> None:
"""Add a request to the waiting queue."""
async with self.lock:
self.waiting_requests[request_id] = time.time()
async def remove_waiting(self, request_id: str) -> None:
"""Remove a request from the waiting queue."""
async with self.lock:
self.waiting_requests.pop(request_id, None)
async def start_processing(self, request_id: str, model_name: str = None) -> None:
"""Mark a request as now processing (model loaded)."""
async with self.lock:
self.waiting_requests.pop(request_id, None)
self.current_request_id = request_id
self.model_name = model_name
async def finish_processing(self) -> None:
"""Mark current request as finished."""
async with self.lock:
self.current_request_id = None
async def is_waiting(self, request_id: str) -> bool:
"""Check if a request is in the waiting queue."""
async with self.lock:
return request_id in self.waiting_requests
async def get_wait_time(self, request_id: str) -> float:
"""Get how long a request has been waiting in seconds."""
async with self.lock:
if request_id in self.waiting_requests:
return time.time() - self.waiting_requests[request_id]
return 0.0
async def get_queue_position(self, request_id: str) -> int:
"""Get the position of a request in the queue (1-based)."""
async with self.lock:
keys = list(self.waiting_requests.keys())
try:
return keys.index(request_id) + 1
except ValueError:
return 0
# Global queue manager
queue_manager = QueueManager()
# =============================================================================
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment