feat!: unify prompts and skills into shared templates and registry

notebooks/example.ipynb

@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "b5fe35fe",
    "metadata": {},
    "outputs": [],
@@ -24,9 +24,9 @@
     "\n",
     "from wags_llm.cache import InMemoryCache\n",
     "from wags_llm.client.bedrock import BedrockClaudeJsonClient\n",
-    "from wags_llm.prompts.base import BasePromptTemplate\n",
-    "from wags_llm.prompts.registry import PromptRegistry\n",
+    "from wags_llm.registry.base import Registry\n",
     "from wags_llm.services.structured_task import StructuredTaskRunner\n",
+    "from wags_llm.templates.base import PromptTemplate\n",
     "\n",
     "logging.basicConfig(\n",
     "    stream=sys.stdout,\n",
@@ -46,12 +46,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "id": "ba6504ff",
    "metadata": {},
    "outputs": [],
    "source": [
-    "class MyPrompt(BasePromptTemplate):\n",
+    "class MyPrompt(PromptTemplate):\n",
     "    name = \"mondo_id_classification\"\n",
     "    version = \"v1\"\n",
     "\n",
@@ -112,7 +112,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "id": "5b14b5ce",
    "metadata": {},
    "outputs": [
@@ -133,12 +133,12 @@
     "    profile_name=\"dev-account\",\n",
     ")\n",
     "\n",
-    "registry = PromptRegistry()\n",
+    "registry = Registry()\n",
     "registry.register(MyPrompt())\n",
     "\n",
     "service = StructuredTaskRunner(\n",
     "    client=client,\n",
-    "    prompt_registry=registry,\n",
+    "    registry=registry,\n",
     "    cache=InMemoryCache(),\n",
     ")"
    ]
@@ -227,7 +227,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "wags-llm",
+   "display_name": "wags-llm (3.11.14)",
    "language": "python",
    "name": "python3"
   },

notebooks/skills.ipynb

@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "e9161ef9",
    "metadata": {},
    "outputs": [],
@@ -26,9 +26,9 @@
     "from pydantic import BaseModel, ConfigDict\n",
     "\n",
     "from wags_llm.client.bedrock import BedrockClaudeJsonClient\n",
+    "from wags_llm.registry.base import Registry\n",
     "from wags_llm.services.structured_task import StructuredTaskRunner\n",
-    "from wags_llm.skills.base import BaseSkillTemplate\n",
-    "from wags_llm.skills.registry import SkillRegistry\n",
+    "from wags_llm.templates.skill_template import SkillTemplate\n",
     "\n",
     "logging.basicConfig(\n",
     "    stream=sys.stdout,\n",
@@ -45,7 +45,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "class VariantCurationSkill(BaseSkillTemplate):\n",
+    "class VariantCurationSkill(SkillTemplate):\n",
     "    skill_path = Path(\"skills/variant_curation_0.1.0.md\")\n",
     "\n",
     "    def build_user_prompt(self, payload: Mapping[str, Any]) -> str:\n",
@@ -96,7 +96,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "id": "c44eeb11",
    "metadata": {},
    "outputs": [
@@ -111,7 +111,7 @@
     }
    ],
    "source": [
-    "registry = SkillRegistry()\n",
+    "registry = Registry()\n",
     "registry.register(skill)\n",
     "\n",
     "MODEL_ID = \"us.anthropic.claude-sonnet-4-6\"\n",
@@ -129,13 +129,13 @@
     "\n",
     "task_runner = StructuredTaskRunner(\n",
     "    client=llm_client,\n",
-    "    skill_registry=registry,\n",
+    "    registry=registry,\n",
     ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "id": "d045cf30",
    "metadata": {},
    "outputs": [
@@ -201,7 +201,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "wags-llm (3.13.5)",
+   "display_name": "wags-llm (3.11.14)",
    "language": "python",
    "name": "python3"
   },
@@ -215,7 +215,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.5"
+   "version": "3.11.14"
   }
  },
  "nbformat": 4,

src/wags_llm/registry/__init__.py

@@ -0,0 +1,11 @@
+"""Registry module.
+
+Store and retrieve versioned prompt and skill templates.
+"""
+
+from wags_llm.registry.base import Registry, build_empty_registry
+
+__all__ = [
+    "Registry",
+    "build_empty_registry",
+]

src/wags_llm/registry/base.py

@@ -0,0 +1,97 @@
+"""Registry.
+
+Maps (name, version, TemplateType) -> template instance.
+Template instances can be either prompts or skills.
+
+Users typically:
+* create prompts or skills in their project
+* register them here or pass a custom registry
+"""
+
+import logging
+from types import MappingProxyType
+
+from wags_llm.templates.base import TemplateType
+from wags_llm.templates.prompt_template import PromptTemplate
+from wags_llm.templates.skill_template import SkillTemplate
+
+_logger = logging.getLogger(__name__)
+
+_TEMPLATE_CLASS_TO_TYPE = MappingProxyType(
+    {
+        SkillTemplate: TemplateType.SKILL,
+        PromptTemplate: TemplateType.PROMPT,
+    }
+)
+
+
+class Registry:
+    """Store and retrieve prompt and skill templates."""
+
+    def __init__(self) -> None:
+        """Initialize an empty template registry."""
+        self._templates: dict[
+            tuple[str, str, TemplateType], PromptTemplate | SkillTemplate
+        ] = {}
+
+    def register(self, template: PromptTemplate | SkillTemplate) -> None:
+        """Register a template.
+
+        :param template: Template instance to register.
+        :raise TypeError: If the template type is unsupported.
+        :raise ValueError: If a template with the same name, version, and template type is already registered.
+        """
+        for cls, mapped_type in _TEMPLATE_CLASS_TO_TYPE.items():
+            if isinstance(template, cls):
+                template_type = mapped_type
+                break
+        else:
+            msg = f"Unsupported template type: {type(template)}"
+            raise TypeError(msg)
+
+        key = (template.name, template.version, template_type)
+
+        _logger.debug(
+            "Registering template: name='%s', version='%s', template_type='%s'",
+            template.name,
+            template.version,
+            template_type.value,
+        )
+
+        if key in self._templates:
+            msg = f"Template already registered:({template.name}, {template.version}, {template_type.value})"
+            _logger.error(msg)
+            raise ValueError(msg)
+
+        self._templates[key] = template
+
+    def get(
+        self,
+        name: str,
+        version: str,
+        template_type: TemplateType,
+    ) -> PromptTemplate | SkillTemplate:
+        """Retrieve a template by name and version.
+
+        :param name: Template name.
+        :param version: Template version.
+        :param template_type: Template type.
+        :return: Registered template.
+        :raise KeyError: If template not found.
+        """
+        key = (name, version, template_type)
+
+        try:
+            return self._templates[key]
+        except KeyError as exc:
+            msg = f"Template not found: ({name}, {version}, {template_type.value})"
+            _logger.exception(msg)
+            raise KeyError(msg) from exc
+
+
+def build_empty_registry() -> Registry:
+    """Create an empty registry.
+
+    :return: New Registry instance.
+    """
+    return Registry()