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")
if args.output is None:
(root / "crates/haematite-cli/embedded/plugin/skill/SKILL.md").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())