haematite-cli 0.21.0

#!/usr/bin/env python3
"""Generate the skill inventory from the served Rust registry without compiling it."""

import argparse
import json
from pathlib import Path
import re
import sys


class GenerationError(Exception):
    pass


def inventory(source):
    declaration = re.search(r"pub fn tools\(\) -> \[ToolSchema; (\d+)\]", source)
    if declaration is None:
        raise GenerationError("skill_registry_declaration_missing")
    quoted = r'"(?:\\.|[^"\\])*"'
    pattern = re.compile(
        r"definition::<([A-Za-z_][A-Za-z0-9_]*),\s*([A-Za-z_][A-Za-z0-9_]*)>\(\s*("
        + quoted + r"),\s*(" + quoted + r"),\s*\)", re.MULTILINE,
    )
    definitions = [(json.loads(name), json.loads(description), input_type, output_type)
                   for input_type, output_type, name, description in pattern.findall(source)]
    if len(definitions) != int(declaration.group(1)):
        raise GenerationError("skill_registry_definition_count_mismatch")
    names = [name for name, description, input_type, output_type in definitions]
    if len(names) != len(set(names)):
        raise GenerationError("skill_registry_duplicate_tool")
    if any(re.fullmatch(r"[a-z][a-z0-9_]*", name) is None for name in names):
        raise GenerationError("skill_registry_tool_name_invalid")
    return definitions


def cell(text):
    return (text.replace("\\", "\\\\").replace("`", "\\`").replace("|", "\\|")
            .replace("\n", " ").replace("\r", " ").replace("\t", " "))


def struct_fields(source, type_name):
    declaration = re.search(
        r"pub struct " + re.escape(type_name) + r"(?:<[^>]*>)?\s*\{(.*?)\n\}",
        source, re.DOTALL,
    )
    if declaration is None:
        raise GenerationError("skill_query_type_missing")
    before = source[:declaration.start()].rsplit("\n\n", 1)[-1]
    if "rename_all" in before:
        raise GenerationError("skill_query_schema_rename_unsupported")
    return declaration.group(1)


def field_name(attributes, name):
    if "schemars" in attributes:
        raise GenerationError("skill_query_schema_rename_unsupported")
    if "rename" in attributes:
        renamed = re.search(r'rename\s*=\s*("(?:\\.|[^"\\])*")', attributes)
        if renamed is None:
            raise GenerationError("skill_query_schema_rename_unsupported")
        name = json.loads(renamed.group(1))
    return name


def cursor_field(source, type_name):
    fields = re.findall(
        r"((?:#\[[^\n]*\]\s*)*)pub\s+([A-Za-z_][A-Za-z0-9_]*)\s*:\s*Option<String>\s*,",
        struct_fields(source, type_name),
    )
    if len(fields) != 1:
        raise GenerationError("skill_query_cursor_invalid")
    return field_name(*fields[0])


def page_cursor_field(source, type_name):
    fields = re.findall(
        r"((?:#\[[^\n]*\]\s*)*)pub\s+([A-Za-z_][A-Za-z0-9_]*)\s*:\s*Page<[A-Za-z_][A-Za-z0-9_]*>\s*,",
        struct_fields(source, type_name),
    )
    if len(fields) != 1:
        raise GenerationError("skill_query_cursor_invalid")
    if re.search(r"pub\s+entries\s*:\s*Vec<T>\s*,", struct_fields(source, "Page")) is None:
        raise GenerationError("skill_query_cursor_invalid")
    return field_name(*fields[0]) + "." + cursor_field(source, "Page")


def render(preamble, definitions, types_source):
    queries = [row for row in definitions if row[0] == "run_query"]
    if len(queries) != 1:
        raise GenerationError("skill_query_tool_missing")
    query = queries[0]
    input_field = cursor_field(types_source, query[2])
    output_field = page_cursor_field(types_source, query[3])
    marker = "{{query_continuation}}"
    if preamble.count(marker) != 1:
        raise GenerationError("skill_query_marker_invalid")
    continuation = (f"For `run_query`, pass output `{cell(output_field)}` back as input "
                    f"`{cell(input_field)}` until a bounded read is complete.")
    page = preamble.rstrip().replace(marker, continuation) + "\n\n| Tool | Purpose |\n| --- | --- |\n"
    return page + "".join(f"| `{cell(name)}` | {cell(description)} |\n"
                          for name, description, input_type, output_type in definitions)


def main():
    root = Path(__file__).resolve().parents[5]
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()
    skill = root / "plugins/haematite/skills/haematite"
    output = args.output if args.output is not None else skill / "SKILL.md"
    try:
        definitions = inventory((Path(__file__).parent.parent / "schema.rs").read_text())
        types_source = (Path(__file__).parent.parent / "types.rs").read_text()
        page = render((skill / "preamble.md").read_text(), definitions, types_source)
        output.write_text(page, encoding="utf-8")
    except (GenerationError, OSError, UnicodeError, json.JSONDecodeError) as error:
        print(f"skill_generation_failed: {error}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())