mirror of https://github.com/razor-ai/soup.git
293 lines
9.9 KiB
Python
293 lines
9.9 KiB
Python
"""Ollama integration utilities — detect, deploy, list, remove models."""
|
|
|
|
import re
|
|
import subprocess
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
# Ollama TEMPLATE blocks for common chat formats
|
|
OLLAMA_TEMPLATES: Dict[str, str] = {
|
|
"chatml": (
|
|
'{{ if .System }}<|im_start|>system\n'
|
|
'{{ .System }}<|im_end|>\n'
|
|
'{{ end }}{{ range .Messages }}'
|
|
'{{ if eq .Role "user" }}<|im_start|>user\n'
|
|
'{{ .Content }}<|im_end|>\n'
|
|
'{{ else if eq .Role "assistant" }}<|im_start|>assistant\n'
|
|
'{{ .Content }}<|im_end|>\n'
|
|
'{{ end }}{{ end }}<|im_start|>assistant\n'
|
|
),
|
|
"llama": (
|
|
'<|begin_of_text|>'
|
|
'{{ if .System }}<|start_header_id|>system<|end_header_id|>\n\n'
|
|
'{{ .System }}<|eot_id|>{{ end }}'
|
|
'{{ range .Messages }}'
|
|
'{{ if eq .Role "user" }}<|start_header_id|>user<|end_header_id|>\n\n'
|
|
'{{ .Content }}<|eot_id|>'
|
|
'{{ else if eq .Role "assistant" }}<|start_header_id|>assistant<|end_header_id|>\n\n'
|
|
'{{ .Content }}<|eot_id|>'
|
|
'{{ end }}{{ end }}'
|
|
'<|start_header_id|>assistant<|end_header_id|>\n\n'
|
|
),
|
|
"mistral": (
|
|
'{{ if .System }}[INST] {{ .System }} [/INST]\n{{ end }}'
|
|
'{{ range .Messages }}'
|
|
'{{ if eq .Role "user" }}[INST] {{ .Content }} [/INST]\n'
|
|
'{{ else if eq .Role "assistant" }}{{ .Content }}</s>\n'
|
|
'{{ end }}{{ end }}'
|
|
),
|
|
"vicuna": (
|
|
'{{ if .System }}{{ .System }}\n\n{{ end }}'
|
|
'{{ range .Messages }}'
|
|
'{{ if eq .Role "user" }}USER: {{ .Content }}\n'
|
|
'{{ else if eq .Role "assistant" }}ASSISTANT: {{ .Content }}</s>\n'
|
|
'{{ end }}{{ end }}ASSISTANT:'
|
|
),
|
|
"zephyr": (
|
|
'{{ if .System }}<|system|>\n{{ .System }}</s>\n{{ end }}'
|
|
'{{ range .Messages }}'
|
|
'{{ if eq .Role "user" }}<|user|>\n{{ .Content }}</s>\n'
|
|
'{{ else if eq .Role "assistant" }}<|assistant|>\n{{ .Content }}</s>\n'
|
|
'{{ end }}{{ end }}<|assistant|>\n'
|
|
),
|
|
}
|
|
|
|
# Map Soup data formats to Ollama template names
|
|
FORMAT_TO_TEMPLATE: Dict[str, str] = {
|
|
"chatml": "chatml",
|
|
"alpaca": "llama",
|
|
"llama": "llama",
|
|
"mistral": "mistral",
|
|
"vicuna": "vicuna",
|
|
"zephyr": "zephyr",
|
|
"sharegpt": "chatml",
|
|
}
|
|
|
|
# Valid model name pattern: alphanumeric, hyphens, underscores, colons (for tags)
|
|
_MODEL_NAME_RE = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9._:-]*$")
|
|
|
|
# Soup-deployed model prefix
|
|
SOUP_MODEL_PREFIX = "soup-"
|
|
|
|
# Allowed Ollama PARAMETER keys (prevents injection of arbitrary directives)
|
|
ALLOWED_OLLAMA_PARAMS = frozenset({
|
|
"temperature", "top_p", "top_k", "num_ctx", "num_predict",
|
|
"stop", "repeat_penalty", "repeat_last_n", "seed", "mirostat",
|
|
"mirostat_tau", "mirostat_eta", "num_gpu", "num_thread",
|
|
"tfs_z", "typical_p", "penalize_newline",
|
|
})
|
|
|
|
|
|
def validate_model_name(name: str) -> Tuple[bool, str]:
|
|
"""Validate Ollama model name.
|
|
|
|
Returns (is_valid, error_message).
|
|
"""
|
|
if not name:
|
|
return False, "Model name cannot be empty"
|
|
if len(name) > 128:
|
|
return False, "Model name too long (max 128 characters)"
|
|
# Block path separators and null bytes
|
|
if "/" in name or "\\" in name or "\0" in name:
|
|
return False, "Model name must not contain path separators or null bytes"
|
|
if not _MODEL_NAME_RE.match(name):
|
|
return False, (
|
|
"Model name must start with alphanumeric and contain only "
|
|
"alphanumeric, hyphens, underscores, dots, or colons"
|
|
)
|
|
return True, ""
|
|
|
|
|
|
def validate_gguf_path(gguf_path: Path) -> Tuple[bool, str]:
|
|
"""Validate GGUF file path — must exist and stay under cwd.
|
|
|
|
Returns (is_valid, error_message).
|
|
"""
|
|
if not gguf_path.exists():
|
|
return False, f"GGUF file not found: {gguf_path}"
|
|
if not gguf_path.is_file():
|
|
return False, f"GGUF path is not a file: {gguf_path}"
|
|
if gguf_path.suffix.lower() != ".gguf":
|
|
return False, "File must have a .gguf extension"
|
|
# Path traversal protection: resolve and check it's under cwd
|
|
resolved = gguf_path.resolve()
|
|
cwd = Path.cwd().resolve()
|
|
try:
|
|
resolved.relative_to(cwd)
|
|
except ValueError:
|
|
return False, "GGUF path must be under the current working directory"
|
|
return True, ""
|
|
|
|
|
|
def detect_ollama() -> Optional[str]:
|
|
"""Check if Ollama is installed, return version string or None."""
|
|
try:
|
|
result = subprocess.run(
|
|
["ollama", "--version"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=10,
|
|
)
|
|
if result.returncode == 0:
|
|
output = result.stdout.strip() or result.stderr.strip()
|
|
# Extract version from output like "ollama version is 0.6.2"
|
|
match = re.search(r"(\d+\.\d+\.\d+)", output)
|
|
if match:
|
|
return match.group(1)
|
|
return output if output else "unknown"
|
|
return None
|
|
except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
|
|
return None
|
|
|
|
|
|
def infer_chat_template(config_format: Optional[str]) -> Optional[str]:
|
|
"""Map a Soup config format to an Ollama template name.
|
|
|
|
Returns template name (key into OLLAMA_TEMPLATES) or None if unknown.
|
|
"""
|
|
if not config_format:
|
|
return None
|
|
return FORMAT_TO_TEMPLATE.get(config_format.lower())
|
|
|
|
|
|
def create_modelfile(
|
|
gguf_path: Path,
|
|
template: Optional[str] = None,
|
|
system_prompt: Optional[str] = None,
|
|
parameters: Optional[Dict[str, str]] = None,
|
|
) -> str:
|
|
"""Generate Ollama Modelfile content.
|
|
|
|
Args:
|
|
gguf_path: Path to the GGUF model file.
|
|
template: Template name (key in OLLAMA_TEMPLATES) or raw template string.
|
|
system_prompt: Optional system prompt.
|
|
parameters: Optional dict of PARAMETER key=value pairs.
|
|
|
|
Returns:
|
|
Modelfile content as string.
|
|
"""
|
|
lines = [f"FROM {gguf_path}"]
|
|
|
|
# Template
|
|
if template:
|
|
template_content = OLLAMA_TEMPLATES.get(template, template)
|
|
lines.append(f'TEMPLATE """{template_content}"""')
|
|
|
|
# System prompt
|
|
if system_prompt:
|
|
# Escape double quotes and strip newlines in system prompt
|
|
escaped = system_prompt.replace('"', '\\"')
|
|
escaped = escaped.replace("\n", " ").replace("\r", "")
|
|
lines.append(f'SYSTEM "{escaped}"')
|
|
|
|
# Parameters — validated against allowlist, no control chars
|
|
if parameters:
|
|
for key, value in sorted(parameters.items()):
|
|
if key not in ALLOWED_OLLAMA_PARAMS:
|
|
raise ValueError(
|
|
f"Unknown Ollama parameter: {key!r}. "
|
|
f"Allowed: {', '.join(sorted(ALLOWED_OLLAMA_PARAMS))}"
|
|
)
|
|
if "\n" in value or "\r" in value or "\0" in value:
|
|
raise ValueError(
|
|
f"Parameter value for {key!r} contains "
|
|
"illegal characters (newline or null)"
|
|
)
|
|
lines.append(f"PARAMETER {key} {value}")
|
|
|
|
return "\n\n".join(lines) + "\n"
|
|
|
|
|
|
def deploy_to_ollama(name: str, modelfile_content: str) -> Tuple[bool, str]:
|
|
"""Deploy a model to Ollama via `ollama create`.
|
|
|
|
Args:
|
|
name: Model name for Ollama.
|
|
modelfile_content: Full Modelfile content.
|
|
|
|
Returns:
|
|
(success, message) tuple.
|
|
"""
|
|
import tempfile
|
|
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="w", suffix=".Modelfile", delete=False, encoding="utf-8"
|
|
) as tmp:
|
|
tmp.write(modelfile_content)
|
|
tmp_path = tmp.name
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
["ollama", "create", name, "-f", tmp_path],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=300,
|
|
)
|
|
if result.returncode == 0:
|
|
return True, result.stdout.strip() or "Model created successfully"
|
|
return False, result.stderr.strip() or "ollama create failed"
|
|
except FileNotFoundError:
|
|
return False, "Ollama not found. Install from https://ollama.com"
|
|
except subprocess.TimeoutExpired:
|
|
return False, "ollama create timed out (5 minutes)"
|
|
except OSError as exc:
|
|
return False, f"Failed to run ollama: {exc}"
|
|
finally:
|
|
Path(tmp_path).unlink(missing_ok=True)
|
|
|
|
|
|
def list_soup_models() -> List[Dict[str, str]]:
|
|
"""List Ollama models deployed by Soup (prefixed with soup-).
|
|
|
|
Returns list of dicts with 'name' and 'size' keys.
|
|
"""
|
|
try:
|
|
result = subprocess.run(
|
|
["ollama", "list"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=10,
|
|
)
|
|
if result.returncode != 0:
|
|
return []
|
|
except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
|
|
return []
|
|
|
|
models = []
|
|
for line in result.stdout.strip().splitlines()[1:]: # Skip header
|
|
parts = line.split()
|
|
if not parts:
|
|
continue
|
|
model_name = parts[0]
|
|
# Filter to soup- prefixed models
|
|
if model_name.startswith(SOUP_MODEL_PREFIX) or (
|
|
":" in model_name and model_name.split(":")[0].startswith(SOUP_MODEL_PREFIX)
|
|
):
|
|
size = parts[2] + " " + parts[3] if len(parts) >= 4 else "unknown"
|
|
models.append({"name": model_name, "size": size})
|
|
return models
|
|
|
|
|
|
def remove_model(name: str) -> Tuple[bool, str]:
|
|
"""Remove a model from Ollama.
|
|
|
|
Returns (success, message) tuple.
|
|
"""
|
|
try:
|
|
result = subprocess.run(
|
|
["ollama", "rm", name],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=30,
|
|
)
|
|
if result.returncode == 0:
|
|
return True, result.stdout.strip() or f"Model '{name}' removed"
|
|
return False, result.stderr.strip() or f"Failed to remove '{name}'"
|
|
except FileNotFoundError:
|
|
return False, "Ollama not found. Install from https://ollama.com"
|
|
except subprocess.TimeoutExpired:
|
|
return False, "ollama rm timed out"
|
|
except OSError as exc:
|
|
return False, f"Failed to run ollama: {exc}"
|