Coverage for src / osiris_cli / provider_registry.py: 0%

1import json 

2from datetime import datetime, timezone 

3from pathlib import Path 

4from typing import Dict, Any, Optional 

5 

6import httpx 

7from openai import OpenAI 

8 

9from .config import CONFIG_DIR, settings, KNOWN_PROVIDERS 

10 

11REGISTRY_FILE = CONFIG_DIR / "providers.json" 

12OPENROUTER_MODELS_URL = "https://openrouter.ai/api/v1/models" 

13 

14 

15def _now_iso() -> str: 

16 return datetime.now(timezone.utc).isoformat() 

17 

18 

19def load_registry() -> Dict[str, Any]: 

20 if REGISTRY_FILE.exists(): 

21 try: 

22 return json.loads(REGISTRY_FILE.read_text()) 

23 except Exception: 

24 return {} 

25 return {} 

26 

27 

28def save_registry(data: Dict[str, Any]) -> None: 

29 REGISTRY_FILE.parent.mkdir(parents=True, exist_ok=True) 

30 REGISTRY_FILE.write_text(json.dumps(data, indent=2)) 

31 

32 

33def refresh_registry(fetch_openrouter: bool = True, provider: Optional[str] = None) -> Dict[str, Any]: 

34 """ 

35 Refresh provider metadata (primarily models) and persist cache. 

36 - OpenRouter: public model listing with pricing (no API key required) 

37 - Others: use OpenAI-compatible /models with configured API key 

38 """ 

39 registry = load_registry() 

40 registry.setdefault("sources", {}) 

41 

42 # OpenRouter public listing 

43 if fetch_openrouter and (provider is None or provider == "openrouter"): 

44 try: 

45 with httpx.Client(timeout=15) as client: 

46 resp = client.get(OPENROUTER_MODELS_URL) 

47 resp.raise_for_status() 

48 data = resp.json().get("data", []) 

49 models = [] 

50 for item in data: 

51 mid = item.get("id") 

52 pricing = item.get("pricing", {}) 

53 models.append( 

54 { 

55 "id": mid, 

56 "prompt": pricing.get("prompt"), 

57 "completion": pricing.get("completion"), 

58 "context_length": item.get("context_length"), 

59 "tags": ["FREE"] if (mid and mid.endswith(":free")) or (pricing.get("prompt") == 0 and pricing.get("completion") == 0) else [], 

60 } 

61 ) 

62 registry["sources"]["openrouter"] = { 

63 "fetched_at": _now_iso(), 

64 "count": len(models), 

65 "models": models, 

66 } 

67 except Exception as e: 

68 registry["sources"]["openrouter"] = { 

69 "fetched_at": _now_iso(), 

70 "error": str(e), 

71 } 

72 

73 # Specific provider refresh (OpenAI-compatible) 

74 if provider and provider != "openrouter": 

75 meta = KNOWN_PROVIDERS.get(provider) 

76 if meta: 

77 key_name = meta.get("api_key_name") 

78 api_key = None 

79 if key_name: 

80 val = getattr(settings, key_name, None) 

81 if val: 

82 api_key = val.get_secret_value() if hasattr(val, "get_secret_value") else val 

83 else: 

84 api_key = "ollama" 

85 try: 

86 client = OpenAI(base_url=meta["base_url"], api_key=api_key or "missing") 

87 models = client.models.list() 

88 items = [{"id": m.id, "tags": []} for m in models.data] 

89 registry["sources"][provider] = { 

90 "fetched_at": _now_iso(), 

91 "count": len(items), 

92 "models": items, 

93 } 

94 except Exception as e: 

95 registry["sources"][provider] = { 

96 "fetched_at": _now_iso(), 

97 "error": str(e), 

98 } 

99 

100 registry["last_refreshed"] = _now_iso() 

101 save_registry(registry) 

102 return registry 

103 

104 

105 

106def sync_all_providers() -> Dict[str, Any]: 

107 """Fetch and cache models for all known providers using configured keys.""" 

108 registry = load_registry() 

109 for prov in KNOWN_PROVIDERS.keys(): 

110 try: 

111 refresh_registry(provider=prov) 

112 except Exception as e: 

113 # Store per-provider error without failing the whole sync 

114 if "sources" not in registry: 

115 registry["sources"] = {} 

116 registry["sources"][prov] = {"fetched_at": _now_iso(), "error": str(e)} 

117 registry["last_synced"] = _now_iso() 

118 save_registry(registry) 

119 return registry 

120 

121 

122def validate_model_cached(provider: str, model_id: str) -> Dict[str, Any]: 

123 """Check whether a model id exists in cached registry for provider.""" 

124 cached = list_cached_models(provider) 

125 models = cached.get("models") or [] 

126 present = False 

127 for m in models: 

128 mid = m.get("id") if isinstance(m, dict) else str(m) 

129 if mid == model_id: 

130 present = True 

131 break 

132 return {"provider": provider, "model": model_id, "present": present} 

133 

134 

135 

136def list_cached_models(source: str = "openrouter") -> Dict[str, Any]: 

137 """Return cached models for a source if present.""" 

138 data = load_registry() 

139 src = data.get("sources", {}).get(source) 

140 if not src: 

141 return {"source": source, "error": "no cached data"} 

142 return {"source": source, "count": src.get("count", 0), "models": src.get("models", [])} 

143 

144 

145def list_free_models(source: str = "openrouter") -> Dict[str, Any]: 

146 """Return only models tagged FREE from the cache for a source.""" 

147 cached = list_cached_models(source) 

148 models = cached.get("models") or [] 

149 free = [m for m in models if isinstance(m, dict) and "FREE" in (m.get("tags") or [])] 

150 return {"source": source, "count": len(free), "models": free} 

151 

152 

153 

154def list_paid_models(source: str = "openrouter") -> Dict[str, Any]: 

155 """Return only models not tagged FREE from the cache for a source.""" 

156 cached = list_cached_models(source) 

157 models = cached.get("models") or [] 

158 paid = [m for m in models if not (isinstance(m, dict) and "FREE" in (m.get("tags") or []))] 

159 return {"source": source, "count": len(paid), "models": paid} 

160 

161 

162def search_cached_models(source: str, pattern: str) -> Dict[str, Any]: 

163 """Search cached models by substring match on id.""" 

164 cached = list_cached_models(source) 

165 models = cached.get("models") or [] 

166 pat = (pattern or "").lower() 

167 matches = [] 

168 for m in models: 

169 mid = m.get("id") if isinstance(m, dict) else str(m) 

170 if pat in mid.lower(): 

171 matches.append(m) 

172 return {"source": source, "count": len(matches), "models": matches} 

173 

174 

175def health_check(provider: Optional[str] = None) -> Dict[str, Any]: 

176 """ 

177 Quick health check by listing models via provider base_url + API key. 

178 """ 

179 prov = provider or settings.provider 

180 meta = KNOWN_PROVIDERS.get(prov) 

181 if not meta: 

182 return {"provider": prov, "status": "unknown-provider"} 

183 

184 key_name = meta.get("api_key_name") 

185 api_key = None 

186 if key_name: 

187 val = getattr(settings, key_name, None) 

188 if val: 

189 api_key = val.get_secret_value() if hasattr(val, "get_secret_value") else val 

190 else: 

191 api_key = "ollama" 

192 

193 client = OpenAI(base_url=meta["base_url"], api_key=api_key) 

194 try: 

195 models = client.models.list() 

196 return { 

197 "provider": prov, 

198 "status": "ok", 

199 "model_count": len(models.data), 

200 } 

201 except Exception as e: 

202 return {"provider": prov, "status": "error", "error": str(e)}