!/usr/bin/env python3
"""
Model Gate — Icarus Staging Model Restriction
Enforces 3B parameter limit on staging environment
"""
import os
import sys
from typing import List, Optional
Allowed models for Icarus staging (3B parameters or smaller)
ALLOWED_MODELS: List[str] = [
"qwen2.5:3b",
"llama3.2:3b",
"phi3:mini",
"phi3:3.8b",
]
Blocked model patterns (anything larger than 4B)
BLOCKED_PATTERNS = [":7b", ":8b", ":13b", ":70b", ":34b", "-coder:", "-instruct:14b"]
def is_model_allowed(model_name: str) -> bool:
"""Check if model is allowed for Icarus staging."""
model_lower = model_name.lower()
# Exact match whitelist
if model_lower in [m.lower() for m in ALLOWED_MODELS]:
return True
# Pattern blacklist
for pattern in BLOCKED_PATTERNS:
if pattern in model_lower:
return False
# Default allow for unknown small models (e.g., custom 3B models)
# But require explicit approval for anything suspicious
if any(size in model_lower for size in [":1b", ":2b", ":3b", ":4b"]):
return True
return False
def get_allowed_model(preferred: Optional[str] = None) -> str:
"""
Get model for Icarus staging.
If preferred is specified and allowed, use it.
Otherwise fall back to first allowed model.
"""
if preferred and is_model_allowed(preferred):
return preferred
if preferred and not is_model_allowed(preferred):
print(f"⚠️ Model '{preferred}' blocked in staging. Using fallback.", file=sys.stderr)
return ALLOWED_MODELS[0] # Default to qwen2.5:3b
def validate_ollama_request(model: str) -> None:
"""
Validate and raise if model is not allowed.
Call this before any Ollama API call in Icarus staging.
"""
if os.environ.get("ICARUS_ENV") != "staging":
return # Only enforce in staging
if not is_model_allowed(model):
raise ValueError(
f"🚫 Model '{model}' is BLOCKED in Icarus staging.\n"
f"Allowed models: {', '.join(ALLOWED_MODELS)}\n"
f"Rationale: Staging simulates M4 Mac mini (3B target)"
)
if name == "main":
# CLI test
test_models = sys.argv[1:] if len(sys.argv) > 1 else [
"qwen2.5:3b",
"llama3.2:3b",
"qwen2.5-coder:7b",
"llama3.1:8b",
"unknown:5b"
]
print("Icarus Staging Model Gate")
print("=" * 40)
for model in test_models:
allowed = "✅ ALLOWED" if is_model_allowed(model) else "🚫 BLOCKED"
print(f"{model:20} {allowed}")