"""Offline compounds, open-attachment fragments and the full named-reaction index."""
from __future__ import annotations
import json
from functools import lru_cache
from importlib.resources import files
from typing import Any, Literal
from pydantic import BaseModel, ConfigDict
from vibesketch.document.derive import Derivation
from vibesketch.document.formats import _merged, import_into
from vibesketch.document.model import Document
from vibesketch.document.reaction_library import (
ReactionRefused,
ReactionSource,
_normalise,
list_reactions,
reaction_fragment,
)
TemplateKind = Literal["compound", "fragment", "example", "scaffold", "reference"]
[docs]
class TemplateEntry(BaseModel):
model_config = ConfigDict(frozen=True)
id: str
name: str
aliases: tuple[str, ...]
category: str
tags: tuple[str, ...] = ()
kind: TemplateKind
sources: tuple[ReactionSource, ...]
note: str
smiles: str | None = None
formula: str | None = None
attachment_points: int = 0
scope: str = "chemical"
@lru_cache(maxsize=1)
def _entries() -> tuple[TemplateEntry, ...]:
raw = files("vibesketch.document").joinpath("data/chemical_templates.json").read_text()
chemical = [TemplateEntry.model_validate(e) for e in json.loads(raw)["entries"]]
reactions = [
TemplateEntry(
id=e.id,
name=e.name,
aliases=e.aliases,
category="Named reactions",
kind=e.kind,
sources=e.sources,
note=e.note,
scope=e.scope,
)
for e in list_reactions()
]
return tuple(sorted([*chemical, *reactions], key=lambda e: e.name.casefold()))
[docs]
def list_templates(query: str = "", *, category: str = "", kind: str = "") -> list[TemplateEntry]:
"""Search names, aliases, formulae and categories; optional exact category/kind filters."""
needle = _normalise(query)
return [
e
for e in _entries()
if (not category or category in (e.category, *e.tags))
and (not kind or e.kind == kind)
and any(
needle in _normalise(s)
for s in (e.id, e.name, *e.aliases, e.category, *e.tags, e.formula or "")
)
]
[docs]
def get_template(identifier: str) -> TemplateEntry:
for entry in _entries():
if entry.id == identifier:
return entry
raise ReactionRefused(f"Unknown template library entry: {identifier!r}")
[docs]
def template_fragment(identifier: str, *, scaffold: bool = False) -> dict[str, Any]:
"""Return editable KET; incomplete reactions require explicit scaffold selection."""
entry = get_template(identifier)
if entry.kind not in ("compound", "fragment"):
return reaction_fragment(identifier, kind="scaffold" if scaffold else "example")
if scaffold:
raise ReactionRefused("Scaffold applies only to named reactions")
assert entry.smiles is not None
ket: dict[str, Any] = dict(import_into(None, entry.smiles, "smiles").document.canvas.data)
if entry.kind == "fragment":
# Unnumbered R sites cannot bind to another drawing's existing R1 definitions.
# KET query atoms keep the whole fragment out of concrete calculation species.
for key, molecule in ket.items():
if key.startswith("mol"):
for atom in molecule["atoms"]:
if atom.get("label") == "*":
atom.pop("label")
atom["type"] = "rg-label"
return ket
[docs]
def insert_template(document: Document, identifier: str, *, scaffold: bool = False) -> Derivation:
"""Append a separate editable block, preserving the document and existing chemistry."""
entry = get_template(identifier)
fragment = template_fragment(identifier, scaffold=scaffold)
merged = _merged(dict(document.canvas.data), fragment, gap=6.0)
result = import_into(document, json.dumps(merged), "ket", mode="replace")
return Derivation(result.document, [entry.note, *result.warnings], result.result)