"""Checked, additive geometry returns from a viewer, without deriving a new drawing.
Callers with a persisted document identity pass it on send and return. Without
one, the exact document digest is the identity, so even unrelated edits refuse.
Atom tokens bind the original ordered graph, including added hydrogens. A viewer
must carry these tokens with reordered atoms; matching elements alone is unsafe.
"""
from __future__ import annotations
import hashlib
import io
import json
import zipfile
from typing import Any
from vibesketch.document.model import Conformer, ConformerSource, Document, Species
from vibesketch.document.qvf import QvfError, _bonds, _modules
from vibesketch.elements import ATOMIC_NUMBER
def _digest(value: Any) -> str:
return (
"sha256:"
+ hashlib.sha256(
json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()
).hexdigest()
)
def _graph(species: Species) -> dict[str, Any]:
if not species.conformers or species.smiles is None:
raise QvfError("geometry return requires a species with a graph and conformer")
c = species.conformers[0]
return {
"symbols": c.symbols,
"bonds": _normal_bonds(_bonds(species, c)),
"smiles": species.smiles,
"charge": species.charge,
"multiplicity": species.multiplicity,
}
def _normal_bonds(bonds: list[dict[str, Any]]) -> list[list[float]]:
result = []
seen = set()
for bond in bonds:
i, j = sorted((int(bond["i"]), int(bond["j"])))
if i == j or (i, j) in seen:
raise QvfError("duplicate or self bond in returned graph")
seen.add((i, j))
result.append([i, j, float(bond["order"])])
return sorted(result)
[docs]
def origin(document: Document, species_id: str, document_id: str | None = None) -> dict[str, Any]:
"""The agreed origin vendor payload; geometry and drawing placement are excluded."""
graph = _graph(document.species_by_id(species_id))
digest = _digest(graph)
result = {
"schema": "vibe-sketch-origin/1",
"document": document_id or _digest(document.to_json_data()),
"revision": document.revision,
"species": species_id,
"conformer": 0,
"atoms": [f"{digest}:{i}" for i in range(len(graph["symbols"]))],
"graph": digest,
}
sequence = document.species_by_id(species_id).sequence
if sequence is not None:
result.update(helm=sequence.helm, helm_inchikey=sequence.inchikey)
return result
[docs]
def return_geometry(document: Document, data: bytes, *, document_id: str | None = None) -> Document:
"""Append a verified conformer, retaining all drawing, graph and calculation fields.
Raises QvfError for missing, ambiguous, foreign or chemically changed mappings.
A stale origin revision is accepted only with the identical graph and atom ids.
"""
from vibesketch.limits import check_archive
_, reader_module = _modules()
try:
check_archive(data)
with reader_module.QvfReader(io.BytesIO(data)) as reader:
if reader.manifest.get("qvf_version") not in (1, 2):
raise QvfError("unsupported QVF version")
origins = [s for s in reader.sections if s.get("kind") == "x_vibesketch.origin"]
structures = [s for s in reader.sections if s.get("kind") == "structure"]
if len(origins) != 1 or len(structures) != 1:
raise QvfError("return needs exactly one origin and one structure")
sent = reader.read_json(origins[0]["members"]["origin"])
section = structures[0]
payload = reader.read_json(section["members"]["structure"])
payload["bonds"] = reader.read_json(section["members"]["bonds"])["pairs"]
source = reader.manifest.get("source") or {}
provenance = reader.manifest.get("provenance") or {}
relaxation = provenance.get("x_vibeview")
return _apply(document, sent, section, payload, source, relaxation, document_id)
except (KeyError, TypeError, ValueError, zipfile.BadZipFile) as error:
raise QvfError(f"geometry return refused: {error}") from None
def _apply(
document: Document,
sent: dict[str, Any],
section: dict[str, Any],
payload: dict[str, Any],
source: dict[str, Any],
relaxation: Any,
document_id: str | None,
) -> Document:
species = document.species_by_id(sent["species"])
expected = origin(document, species.id, document_id)
for key in ("schema", "document", "species", "graph"):
if sent.get(key) != expected[key]:
raise QvfError(f"origin {key} no longer matches the document")
revision = sent.get("revision")
if type(revision) is not int or revision < 0 or revision > document.revision:
raise QvfError("invalid or future origin revision")
ids = sent.get("atoms")
if not isinstance(ids, list) or any(not isinstance(i, str) for i in ids):
raise QvfError("missing atom mapping")
if len(ids) != len(set(ids)) or set(ids) != set(expected["atoms"]):
raise QvfError("missing, duplicate or changed atom mapping")
atoms = payload["atoms"]
if len(atoms) != len(ids) or any(payload.get("pbc") or []):
raise QvfError("changed atom count or periodic structure")
indices = [expected["atoms"].index(token) for token in ids]
ordered = [atoms[ids.index(token)] for token in expected["atoms"]]
graph = _graph(species)
if [a["symbol"] for a in ordered] != graph["symbols"]:
raise QvfError("mapped atom elements changed")
for atom in ordered:
if (
atom.get("atomic_number", ATOMIC_NUMBER[atom["symbol"]])
!= ATOMIC_NUMBER[atom["symbol"]]
):
raise QvfError("mapped atomic number contradicts element")
for key in ("charge", "multiplicity", "smiles"):
if section.get(key) != graph[key]:
raise QvfError(f"returned {key} changed or missing")
bonds = payload.get("bonds")
if not isinstance(bonds, list):
raise QvfError("returned bond graph is missing; distances are not a mapping")
remapped = []
for b in bonds:
if type(b["i"]) is not int or type(b["j"]) is not int:
raise QvfError("invalid bond indices")
if not (0 <= b["i"] < len(ids) and 0 <= b["j"] < len(ids)):
raise QvfError("bond index outside atom mapping")
remapped.append({"i": indices[b["i"]], "j": indices[b["j"]], "order": b["order"]})
if _normal_bonds(remapped) != graph["bonds"]:
raise QvfError("returned bonds or bond orders changed")
conformer = Conformer(
method="returned",
symbols=[a["symbol"] for a in ordered],
xyz_angstrom=[a["position"] for a in ordered],
source=ConformerSource(
program=source.get("program", "unknown"),
version=source.get("version"),
relaxation=relaxation,
),
)
updated = document.to_json_data()
target = next(s for s in updated["species"] if s["id"] == species.id)
target.setdefault("conformers", []).append(conformer.model_dump(exclude_none=True))
return Document.model_validate(updated)