81 lines
2.9 KiB
Python
81 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional
|
|
|
|
from config import MODELS_DIR
|
|
from engine_loader import EngineLoadError, load_engine_class
|
|
from .base import ModelInfo, Provider, ProviderError
|
|
|
|
|
|
class LocalEngineProvider(Provider):
|
|
name = "Local Engine"
|
|
|
|
def __init__(self, search_paths: Optional[List[Path]] = None) -> None:
|
|
try:
|
|
runtime_cls = load_engine_class()
|
|
self.runtime = runtime_cls()
|
|
self.runtime.setup()
|
|
except EngineLoadError as exc:
|
|
raise ProviderError(str(exc)) from exc
|
|
except RuntimeError as exc:
|
|
raise ProviderError(str(exc)) from exc
|
|
|
|
paths: List[Path] = [MODELS_DIR]
|
|
env_override = os.getenv("LOCAL_LLM_PATHS")
|
|
if env_override:
|
|
for entry in env_override.split(os.pathsep):
|
|
entry_clean = entry.strip()
|
|
if not entry_clean:
|
|
continue
|
|
entry_path = Path(entry_clean).expanduser()
|
|
if entry_path not in paths:
|
|
paths.append(entry_path)
|
|
if search_paths:
|
|
for candidate in search_paths:
|
|
if candidate not in paths:
|
|
paths.append(candidate)
|
|
|
|
self.search_paths = paths
|
|
self._model_index: Dict[str, Path] = {}
|
|
|
|
def list_models(self) -> List[ModelInfo]:
|
|
discovered = self.runtime.discover_gguf_models(self.search_paths)
|
|
self._model_index.clear()
|
|
models: List[ModelInfo] = []
|
|
for path in discovered:
|
|
model_id = self._register_model(path)
|
|
models.append(ModelInfo(id=model_id, display_name=path.stem))
|
|
if not models:
|
|
raise ProviderError(
|
|
"No GGUF models were found. Place models inside the 'models' directory or set LOCAL_LLM_PATHS."
|
|
)
|
|
return models
|
|
|
|
def run_prompt(self, model_id: str, prompt: str, *, temperature: float) -> Dict[str, object]:
|
|
try:
|
|
model_path = self._model_index[model_id]
|
|
except KeyError as exc:
|
|
raise ProviderError(f"Model '{model_id}' is not registered. Refresh the model list and try again.") from exc
|
|
|
|
self.runtime.load_model(model_path)
|
|
return self.runtime.generate(prompt, temperature=temperature)
|
|
|
|
def _register_model(self, path: Path) -> str:
|
|
for existing_id, existing_path in self._model_index.items():
|
|
if path == existing_path:
|
|
return existing_id
|
|
|
|
base_id = path.stem
|
|
candidate = base_id
|
|
counter = 1
|
|
while candidate in self._model_index:
|
|
digest = hashlib.sha1(str(path).encode("utf-8")).hexdigest()[:6]
|
|
candidate = f"{base_id}-{digest}-{counter}"
|
|
counter += 1
|
|
|
|
self._model_index[candidate] = path
|
|
return candidate
|