File size: 2,038 Bytes
6d88ccb |
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 63 64 65 66 |
from __future__ import annotations
import importlib
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List
import yaml
REQUIRED_FIELDS = {"name", "type", "version", "description", "capabilities", "entrypoint"}
@dataclass
class CatalogEntry:
name: str
type: str
version: str
description: str
capabilities: List[str]
entrypoint: str
provider: str | None = None
def load(self) -> Any:
module_name, attr = self.entrypoint.rsplit(".", 1)
module = importlib.import_module(module_name)
return getattr(module, attr)
class CatalogRegistry:
def __init__(self, entries: List[CatalogEntry]) -> None:
self.entries = entries
@classmethod
def _load_file(cls, path: Path) -> List[CatalogEntry]:
data = yaml.safe_load(path.read_text()) if path.exists() else []
entries: List[CatalogEntry] = []
for raw in data or []:
missing = REQUIRED_FIELDS - set(raw)
if missing:
raise ValueError(f"Catalog entry missing fields {missing} in {path}")
entries.append(CatalogEntry(**raw))
return entries
@classmethod
def from_default(cls) -> "CatalogRegistry":
base = Path(__file__).parent.parent / "catalogs"
entries: List[CatalogEntry] = []
for name in ["models.yaml", "tools.yaml", "plugins.yaml"]:
entries.extend(cls._load_file(base / name))
return cls(entries)
def find(self, *, type: str | None = None, capability: str | None = None) -> Iterable[CatalogEntry]:
for entry in self.entries:
if type and entry.type != type:
continue
if capability and capability not in entry.capabilities:
continue
yield entry
def list_all(self) -> List[Dict[str, str]]:
return [
{"type": e.type, "name": e.name, "description": e.description, "version": e.version}
for e in self.entries
]
|