wehub-resource-sync
202 lines
7.7 KiB
Python
202 lines
7.7 KiB
Python
import json
|
|
import logging
|
|
import re
|
|
from collections.abc import AsyncGenerator
|
|
from urllib.parse import urlparse
|
|
|
|
import aiohttp
|
|
from fastapi import HTTPException
|
|
|
|
from constants.supported_ollama_models import SUPPORTED_OLLAMA_MODELS
|
|
from models.ollama_model_status import OllamaModelStatus
|
|
from utils.get_env import get_ollama_url_env
|
|
|
|
LOGGER = logging.getLogger(__name__)
|
|
|
|
def _extract_ollama_parameter_suffix(model_ref: str) -> str:
|
|
suffix_match = re.search(
|
|
r":((?:[0-9]+(?:\.[0-9]+)?(?:x[0-9]+(?:\.[0-9]+)?)?)b)\b",
|
|
model_ref,
|
|
re.IGNORECASE,
|
|
)
|
|
if suffix_match:
|
|
return suffix_match.group(1).upper()
|
|
return ""
|
|
|
|
|
|
def _build_ollama_library_models() -> list[dict]:
|
|
models: list[dict] = []
|
|
for model_id, metadata in sorted(SUPPORTED_OLLAMA_MODELS.items()):
|
|
parameters = _extract_ollama_parameter_suffix(model_id)
|
|
models.append(
|
|
{
|
|
"name": metadata.value,
|
|
"description": metadata.label,
|
|
"parameters": parameters if parameters else None,
|
|
"size": metadata.size,
|
|
}
|
|
)
|
|
return models
|
|
|
|
|
|
OLLAMA_LIBRARY_MODELS = _build_ollama_library_models()
|
|
|
|
|
|
def _get_ollama_url(ollama_url: str | None = None) -> str:
|
|
resolved_url = (
|
|
ollama_url or get_ollama_url_env() or "http://localhost:11434"
|
|
).strip()
|
|
if not resolved_url:
|
|
resolved_url = "http://localhost:11434"
|
|
if any(ord(ch) < 32 for ch in resolved_url):
|
|
raise HTTPException(status_code=400, detail="Invalid Ollama URL")
|
|
parsed_url = urlparse(resolved_url)
|
|
if not parsed_url.scheme:
|
|
resolved_url = f"http://{resolved_url}"
|
|
parsed_url = urlparse(resolved_url)
|
|
if parsed_url.scheme not in {"http", "https"} or not parsed_url.netloc:
|
|
raise HTTPException(status_code=400, detail="Invalid Ollama URL")
|
|
return resolved_url.rstrip("/")
|
|
|
|
|
|
def _ollama_unreachable_error(ollama_url: str | None = None) -> HTTPException:
|
|
resolved_ollama_url = _get_ollama_url(ollama_url)
|
|
return HTTPException(
|
|
status_code=503,
|
|
detail=(
|
|
f"Could not connect to Ollama at {resolved_ollama_url}. "
|
|
"Make sure Ollama is running and reachable from Presenton. "
|
|
"When Presenton runs in Docker, use host.docker.internal instead of localhost."
|
|
),
|
|
)
|
|
|
|
|
|
def _extract_ollama_parameter_count(model_name: str, model_details: dict | None = None) -> str:
|
|
details = model_details or {}
|
|
details_parameter_size = details.get("parameter_size")
|
|
if isinstance(details_parameter_size, str) and details_parameter_size.strip():
|
|
return details_parameter_size.strip().upper()
|
|
|
|
parameters_from_name = _extract_ollama_parameter_suffix(model_name)
|
|
if parameters_from_name:
|
|
return parameters_from_name
|
|
|
|
supported_model = SUPPORTED_OLLAMA_MODELS.get(model_name.lower())
|
|
if supported_model:
|
|
return _extract_ollama_parameter_suffix(supported_model.value)
|
|
return ""
|
|
|
|
|
|
async def list_available_ollama_models(
|
|
ollama_url: str | None = None,
|
|
) -> list[OllamaModelStatus]:
|
|
base_url = _get_ollama_url(ollama_url)
|
|
try:
|
|
async with aiohttp.ClientSession(
|
|
timeout=aiohttp.ClientTimeout(total=10)
|
|
) as session:
|
|
async with session.get(
|
|
f"{base_url}/api/tags",
|
|
) as response:
|
|
if response.status == 200:
|
|
pulled_models = await response.json(content_type=None)
|
|
models = (
|
|
pulled_models.get("models")
|
|
if isinstance(pulled_models, dict)
|
|
else None
|
|
)
|
|
if not isinstance(models, list):
|
|
raise HTTPException(
|
|
status_code=502,
|
|
detail="Ollama returned an invalid models response",
|
|
)
|
|
return [
|
|
OllamaModelStatus(
|
|
name=m.get("model") or m.get("name"),
|
|
parameters=_extract_ollama_parameter_count(
|
|
m.get("model") or m.get("name") or "",
|
|
m.get("details") if isinstance(m, dict) else None,
|
|
)
|
|
or None,
|
|
size=m.get("size") or 0,
|
|
status="pulled",
|
|
downloaded=m.get("size") or 0,
|
|
done=True,
|
|
)
|
|
for m in models
|
|
if isinstance(m, dict) and (m.get("model") or m.get("name"))
|
|
]
|
|
elif response.status == 403:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail="Forbidden: Please check your Ollama Configuration",
|
|
)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=response.status,
|
|
detail=f"Failed to list Ollama models: {response.status}",
|
|
)
|
|
except HTTPException:
|
|
raise
|
|
except (aiohttp.ClientError, TimeoutError, json.JSONDecodeError) as error:
|
|
raise _ollama_unreachable_error(ollama_url) from error
|
|
|
|
|
|
def get_ollama_library_models() -> list[dict]:
|
|
return OLLAMA_LIBRARY_MODELS
|
|
|
|
|
|
async def pull_ollama_model(
|
|
model_name: str,
|
|
ollama_url: str | None = None,
|
|
) -> AsyncGenerator[str, None]:
|
|
base_url = _get_ollama_url(ollama_url)
|
|
try:
|
|
async with aiohttp.ClientSession(
|
|
timeout=aiohttp.ClientTimeout(total=None)
|
|
) as session:
|
|
async with session.post(
|
|
f"{base_url}/api/pull",
|
|
json={"name": model_name, "stream": True, "insecure": False},
|
|
) as response:
|
|
if response.status != 200:
|
|
body = await response.text()
|
|
yield f"event: error\ndata: {json.dumps({'detail': body or 'Pull failed'})}\n\n"
|
|
return
|
|
|
|
async for line in response.content:
|
|
decoded = line.decode("utf-8").strip()
|
|
if not decoded:
|
|
continue
|
|
try:
|
|
data = json.loads(decoded)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
if data.get("error"):
|
|
yield f"event: error\ndata: {json.dumps({'detail': data['error']})}\n\n"
|
|
return
|
|
|
|
if data.get("status") == "success":
|
|
yield f"event: response\ndata: {json.dumps({'type': 'complete', 'status': 'success', 'model': model_name})}\n\n"
|
|
return
|
|
|
|
total = data.get("total")
|
|
completed = data.get("completed")
|
|
status = data.get("status", "")
|
|
|
|
if total and completed is not None:
|
|
progress = round((completed / total) * 100, 1)
|
|
yield (
|
|
f"event: response\ndata: "
|
|
f"{json.dumps({'type': 'progress', 'status': status, 'total': total, 'completed': completed, 'progress': progress})}\n\n"
|
|
)
|
|
else:
|
|
yield (
|
|
f"event: response\ndata: "
|
|
f"{json.dumps({'type': 'status', 'status': status})}\n\n"
|
|
)
|
|
except (aiohttp.ClientError, TimeoutError) as error:
|
|
LOGGER.error("Ollama pull error: %s", error)
|
|
yield f"event: error\ndata: {json.dumps({'detail': f'Could not connect to Ollama at {base_url}'})}\n\n"
|