File size: 3,160 Bytes
287f01b b17b915 287f01b b17b915 287f01b b17b915 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
# llm_router.py — enruta llamadas a Spaces remotos según config.yaml
from __future__ import annotations
from typing import Any, Dict, List, Optional
from pathlib import Path
import os
import yaml
from remote_clients import InstructClient, VisionClient, ToolsClient, ASRClient
def load_yaml(path: str) -> Dict[str, Any]:
p = Path(path)
if not p.exists():
return {}
return yaml.safe_load(p.read_text(encoding="utf-8")) or {}
class LLMRouter:
def __init__(self, cfg: Dict[str, Any]):
self.cfg = cfg
self.rem = cfg.get("models", {}).get("routing", {}).get("use_remote_for", [])
base_user = cfg.get("remote_spaces", {}).get("user", "veureu")
eps = cfg.get("remote_spaces", {}).get("endpoints", {})
token_enabled = cfg.get("security", {}).get("use_hf_token", False)
hf_token = os.getenv(cfg.get("security", {}).get("hf_token_env", "HF_TOKEN")) if token_enabled else None
def mk(endpoint_key: str, cls):
info = eps.get(endpoint_key, {})
base_url = info.get("base_url") or f"https://{base_user}-{info.get('space')}.hf.space"
use_gradio = (info.get("client", "gradio") == "gradio")
timeout = int(cfg.get("remote_spaces", {}).get("http", {}).get("timeout_seconds", 180))
return cls(base_url=base_url, use_gradio=use_gradio, hf_token=hf_token, timeout=timeout)
self.clients = {
"salamandra-instruct": mk("salamandra-instruct", InstructClient),
"salamandra-vision": mk("salamandra-vision", VisionClient),
"salamandra-tools": mk("salamandra-tools", ToolsClient),
"whisper-catalan": mk("whisper-catalan", ASRClient),
}
# ---- INSTRUCT ----
def instruct(self, prompt: str, system: Optional[str] = None, model: str = "salamandra-instruct", **kwargs) -> str:
if model in self.rem:
return self.clients[model].generate(prompt, system=system, **kwargs) # type: ignore
raise RuntimeError(f"Modelo local no implementado para: {model}")
# ---- VISION ----
def vision_describe(self, image_paths: List[str], context: Optional[Dict[str, Any]] = None, model: str = "salamandra-vision", **kwargs) -> List[str]:
if model in self.rem:
return self.clients[model].describe(image_paths, context=context, **kwargs) # type: ignore
raise RuntimeError(f"Modelo local no implementado para: {model}")
# ---- TOOLS ----
def chat_with_tools(self, messages: List[Dict[str, str]], tools: Optional[List[Dict[str, Any]]] = None, model: str = "salamandra-tools", **kwargs) -> Dict[str, Any]:
if model in self.rem:
return self.clients[model].chat(messages, tools=tools, **kwargs) # type: ignore
raise RuntimeError(f"Modelo local no implementado para: {model}")
# ---- ASR ----
def asr_transcribe(self, audio_path: str, model: str = "whisper-catalan", **kwargs) -> Dict[str, Any]:
if model in self.rem:
return self.clients[model].transcribe(audio_path, **kwargs) # type: ignore
raise RuntimeError(f"Modelo local no implementado para: {model}")
|