simit 0.1.0__py3-none-any.whl

simit/skills/base.py

@@ -0,0 +1,172 @@
+"""The ``Skill`` interface: how a textual image description becomes an image.
+
+A skill turns a description into a compact specification written by the
+model (``prompt`` -> model -> ``parse``), then draws it deterministically
+(``render``). Parser/renderer errors are fed back to the model for repair.
+Subclass ``Skill`` to add a new visual domain::
+
+    class SheetMusic(simit.Skill):
+        name = "sheet_music"
+        route_hint = "musical notation: staves, notes, chords"
+        examples = [("Draw a C major scale", "NOTES: C4 D4 E4 F4 G4 A4 B4 C5")]
+        def parse(self, text): ...       # raise SpecError(msg); msg is fed back
+        def render(self, spec): ...      # -> PIL.Image
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Optional, Sequence
+
+from PIL import Image
+
+from ..utils import clean_response
+
+if TYPE_CHECKING:
+    from .context import SkillContext
+
+
+class SpecError(ValueError):
+    """Raise from ``Skill.parse``/``Skill.render`` when a specification is
+    unusable. The message is shown to the model in the repair prompt, so make
+    it specific ("no NODE: lines found", "undefined point 'C'")."""
+
+
+@dataclass
+class Realization:
+    image: Image.Image
+    skill: str
+    spec_text: str = ""            # the model output the image was rendered from
+    attempts: int = 1
+    extra: dict = field(default_factory=dict)
+
+
+# Completion-style few-shot framing ("Request: ... / Specification:"): models
+# follow a custom format far more reliably when the prompt ends where the spec
+# starts (BAGEL: 15/15 parseable vs 8/15 with a trailing instruction).
+_DEFAULT_PROMPT = """You are a precise {title} specification writer. Given a description of \
+{hint}, you output ONLY its specification, in exactly the format shown in the examples below \
+-- nothing else (no explanations, no prose).
+
+{examples}{note}Request: {{request}}
+Specification:
+"""
+
+_REPAIR_NOTE = """A previous specification for the next request could not be used ({reason}):
+{previous}
+
+Write a corrected one.
+
+"""
+
+_CUSTOM_PROMPT_REPAIR = """{prompt}
+
+Your previous answer could not be used ({reason}). It was:
+{previous}
+
+Output ONLY a corrected specification, nothing else."""
+
+
+class Skill:
+    """Base class for a rendering skill. See the module docstring."""
+
+    #: Unique identifier; the router answers with ``name.upper()``.
+    name: str = ""
+    #: Category definition shown to the router ("when should I be used?").
+    route_hint: str = ""
+    #: Example requests that belong to this skill (router few-shots).
+    #: Defaults to the requests in ``examples``.
+    route_examples: Sequence[str] = ()
+    #: ``(request, spec_text)`` few-shot pairs for the auto-built spec prompt.
+    examples: Sequence[tuple[str, str]] = ()
+    #: Full spec-prompt template containing ``{request}``; auto-built from
+    #: ``examples`` when left as ``None``.
+    prompt: Optional[str] = None
+    max_new_tokens: int = 512
+    #: Decoding temperature of the first attempt and of repair attempts
+    #: (greedy repairs on a near-identical context repeat the same mistake).
+    first_temperature: float = 0.0
+    retry_temperature: float = 0.5
+    #: Critic rubric: "strict" (every detail must match, for photos) or
+    #: "structured" (right kind of visual, readable, ignore exact values).
+    verify_rubric: str = "structured"
+    #: Render in a worker process (True) or a thread of the main process.
+    #: Process rendering needs the skill to be importable/picklable.
+    render_in_subprocess: bool = True
+
+    # ------------------------------------------------------------------ prompts
+    @property
+    def label(self) -> str:
+        return self.name.upper()
+
+    def build_prompt(self, request: str) -> str:
+        template = self.prompt if self.prompt is not None else self._default_prompt()
+        return template.format(request=request)
+
+    def build_repair_prompt(self, request: str, previous: str, reason: str) -> str:
+        if self.prompt is not None:
+            return _CUSTOM_PROMPT_REPAIR.format(prompt=self.build_prompt(request), previous=previous, reason=reason)
+        note = _REPAIR_NOTE.format(reason=reason, previous=previous.strip())
+        return self._default_prompt(note).format(request=request)
+
+    def _default_prompt(self, note: str = "") -> str:
+        blocks = [f"Request: {req}\nSpecification:\n{spec.strip()}\n\n" for req, spec in self.examples]
+        title = self.name.replace("_", " ")
+        hint = self.route_hint.strip() or f"a {title} image"
+
+        def esc(text):  # user/model text must survive the final .format(request=...)
+            return text.replace("{", "{{").replace("}", "}}")
+        return _DEFAULT_PROMPT.format(title=title, hint=esc(hint), examples=esc("".join(blocks)), note=esc(note))
+
+    def get_route_examples(self) -> list[str]:
+        return list(self.route_examples) or [req for req, _ in self.examples]
+
+    # ------------------------------------------------------------- spec -> image
+    def shortcut(self, request: str) -> Any:
+        """Return a ready spec for ``request`` to skip the model, or None."""
+        return None
+
+    def clean(self, text: str) -> str:
+        return clean_response(text).strip()
+
+    def parse(self, text: str) -> Any:
+        raise NotImplementedError
+
+    def render(self, spec: Any) -> Image.Image:
+        raise NotImplementedError
+
+    # ------------------------------------------------------------- realization
+    async def realize(self, ctx: "SkillContext", request: str, retries: int = 2) -> Optional[Realization]:
+        """prompt -> model -> parse -> render, repairing on errors. Returns
+        None when every attempt failed."""
+        shortcut = self.shortcut(request)
+        if shortcut is not None:
+            try:
+                image = await ctx.render(self, shortcut)
+                return Realization(image, self.name, spec_text=str(shortcut), attempts=0)
+            except SpecError:
+                pass
+        prompt = self.build_prompt(request)
+        for attempt in range(retries + 1):
+            temperature = self.first_temperature if attempt == 0 else self.retry_temperature
+            text = self.clean(await ctx.generate(prompt, max_new_tokens=self.max_new_tokens,
+                                                 temperature=temperature))
+            try:
+                spec = self.parse(text)
+                image = await ctx.render(self, spec)
+            except SpecError as e:
+                ctx.log(f"[{self.name}] attempt {attempt}: {e}")
+                prompt = self.build_repair_prompt(request, text, str(e))
+                continue
+            return Realization(image, self.name, spec_text=text, attempts=attempt + 1)
+        return None
+
+
+_FENCE_RE = re.compile(r"```(?:[\w-]+)?\s*\n(.*?)```", re.DOTALL)
+
+
+def extract_fenced_block(text: str) -> Optional[str]:
+    """Content of the first fenced code block, or None."""
+    m = _FENCE_RE.search(text)
+    return m.group(1).strip() if m else None