Доработка управления моделями

app/templates.py

@@ -1,9 +1,17 @@
-"""Template loading with startup validation of graph links and mappings."""
+"""Model registry loading: read each workflow, validate links/tokens, and
+decide which model an endpoint may run.
+
+Startup validation of a workflow is local here (broken links + token check);
+the collected graphs are then dry-injected and sent to the server elsewhere
+(see main.lifespan).
+"""
 import json
+from dataclasses import dataclass, field
 from pathlib import Path
 
 from .config import Settings
-from .errors import AdapterError
+from .errors import AdapterError, ConfigError, InvalidRequest
+from .injector import extract_tokens
 
 
 class TemplateError(AdapterError):
@@ -13,6 +21,29 @@ class TemplateError(AdapterError):
         super().__init__(500, message, "server_error", "template_error")
 
 
+@dataclass(frozen=True)
+class Model:
+    id: str
+    graph: dict
+    declared: frozenset[str]
+    owned_by: str = "comfy-adapter"
+    workflow: Path = field(default=Path("."))
+
+
+# Token sets belong to the endpoint, not the model. A model is valid for an
+# endpoint iff required(endpoint) <= declared(graph) <= supplied(endpoint).
+ENDPOINT_VOCAB: dict[str, dict[str, frozenset[str]]] = {
+    "generations": {
+        "supplied": frozenset({"prompt", "negative", "seed", "width", "height"}),
+        "required": frozenset({"prompt"}),
+    },
+    "edits": {
+        "supplied": frozenset({"prompt", "negative", "seed", "image"}),
+        "required": frozenset({"prompt", "image"}),
+    },
+}
+
+
 def find_broken_links(graph: dict) -> list[str]:
     """References to node ids that do not exist in the graph."""
     errors = []
@@ -24,39 +55,58 @@ def find_broken_links(graph: dict) -> list[str]:
     return errors
 
 
-def check_mapping(graph: dict, mapping: dict, name: str) -> list[str]:
-    """Every mapped node id must exist with a scalar input field."""
-    errors = []
-    for key, spec in mapping.items():
-        node = graph.get(spec["node"])
-        if node is None:
-            errors.append(f"mapping '{name}.{key}': node '{spec['node']}' not found in template")
-            continue
-        for field in _mapped_fields(spec):
-            if field not in node.get("inputs", {}) or isinstance(node["inputs"][field], list):
-                errors.append(f"mapping '{name}.{key}': node '{spec['node']}' has no scalar input '{field}'")
-    return errors
+def is_compatible(declared: frozenset[str], endpoint: str) -> bool:
+    vocab = ENDPOINT_VOCAB[endpoint]
+    return vocab["required"] <= declared and declared <= vocab["supplied"]
 
 
-def _mapped_fields(spec: dict) -> list[str]:
-    if "width_field" in spec:
-        return [spec["width_field"], spec["height_field"]]
-    return [spec["field"]]
+def select_model(models: dict[str, Model], model: str | None, endpoint: str) -> Model:
+    """Pick the model for a request by id, or the first compatible one.
+
+    An explicit id that is unknown or incompatible with the endpoint is a 400,
+    as is a missing id when no registered model is compatible.
+    """
+    if model is not None:
+        found = models.get(model)
+        if found is None:
+            raise InvalidRequest(f"unknown model '{model}'")
+        if not is_compatible(found.declared, endpoint):
+            raise InvalidRequest(f"model '{model}' is not compatible with this endpoint")
+        return found
+    for candidate in models.values():  # declaration order = priority
+        if is_compatible(candidate.declared, endpoint):
+            return candidate
+    raise InvalidRequest(f"no model in the registry is compatible with this endpoint")
 
 
-def load_template(path: Path, graph_name: str, mapping: dict, mapping_name: str) -> dict:
-    with open(path, encoding="utf-8") as f:
-        graph = json.load(f)
-    errors = find_broken_links(graph) + check_mapping(graph, mapping[mapping_name], mapping_name)
+def _read_graph(path: Path, model_id: str) -> dict:
+    try:
+        with open(path, encoding="utf-8") as f:
+            return json.load(f)
+    except FileNotFoundError as e:
+        raise TemplateError(f"model '{model_id}': workflow file not found: {path}") from e
+    except json.JSONDecodeError as e:
+        raise TemplateError(f"model '{model_id}': workflow is not valid JSON ({path}): {e}") from e
+
+
+def load_model(model_id: str, workflow_path: Path, owned_by: str) -> Model:
+    graph = _read_graph(workflow_path, model_id)
+    errors = find_broken_links(graph)
+    declared, unknown = extract_tokens(graph)
+    errors += unknown
     if errors:
-        raise TemplateError(f"workflow template '{graph_name}' ({path}) is invalid:\n  " + "\n  ".join(errors))
-    return graph
+        raise TemplateError(
+            f"workflow for model '{model_id}' ({workflow_path}) is invalid:\n  " + "\n  ".join(errors)
+        )
+    return Model(id=model_id, graph=graph, declared=frozenset(declared), owned_by=owned_by, workflow=workflow_path)
 
 
-def load_all(settings: Settings) -> tuple[dict, dict]:
-    mapping = settings.mapping()
-    create = load_template(
-        settings.create_template_path, "workflow_create.json", mapping, "create"
-    )
-    edit = load_template(settings.edit_template_path, "workflow_edit.json", mapping, "edit")
-    return create, edit
+def load_registry(settings: Settings) -> dict[str, Model]:
+    """Load every model in the registry, preserving order for default selection."""
+    base_dir = settings.models_file.parent
+    registry = settings.load_registry()
+    models: dict[str, Model] = {}
+    for model_id, entry in registry.items():
+        path = entry.workflow if entry.workflow.is_absolute() else base_dir / entry.workflow
+        models[model_id] = load_model(model_id, path, entry.owned_by)
+    return models