Improve runtime naming of public API types (#74)

Change the names of public API types at schema generation
time instead of aliasing them at import time.

This makes the name reported at runtime match the import
name in the public SDK API.

Closes #73

Author: ncoghlan

sdk-schema/sync-sdk-schema.py

@@ -334,6 +334,36 @@ def _infer_schema_unions() -> None:
     _INFERRED_SCHEMA_PATH.write_text(json.dumps(processed_schema, indent=2))
 
 
+# Unfortunately, "aliases" in the code generator isn't full type renaming
+# Instead, these are handled as part of the AST transformation step
+_DATA_MODEL_NAME_OVERRIDES = {
+    # Prettier chat history type names
+    "ChatMessageData": "AnyChatMessage",
+    "ChatMessageDataUser": "UserMessage",
+    "ChatMessageDataSystem": "SystemPrompt",
+    "ChatMessageDataAssistant": "AssistantResponse",
+    "ChatMessageDataTool": "ToolResultMessage",
+    "ChatMessageDataUserDict": "UserMessageDict",
+    "ChatMessageDataSystemDict": "SystemPromptDict",
+    "ChatMessageDataAssistantDict": "AssistantResponseDict",
+    "ChatMessageDataToolDict": "ToolResultMessageDict",
+    "ChatMessagePartFileData": "FileHandle",
+    "ChatMessagePartFileDataDict": "FileHandleDict",
+    "ChatMessagePartTextData": "TextData",
+    "ChatMessagePartTextDataDict": "TextDataDict",
+    "ChatMessagePartToolCallRequestData": "ToolCallRequestData",
+    "ChatMessagePartToolCallRequestDataDict": "ToolCallRequestDataDict",
+    "ChatMessagePartToolCallResultData": "ToolCallResultData",
+    "ChatMessagePartToolCallResultDataDict": "ToolCallResultDataDict",
+    "FunctionToolCallRequestDict": "ToolCallRequestDict",
+    # Prettier channel creation type names
+    "LlmChannelPredictCreationParameter": "PredictionChannelRequest",
+    "LlmChannelPredictCreationParameterDict": "PredictionChannelRequestDict",
+    "RepositoryChannelDownloadModelCreationParameter": "DownloadModelChannelRequest",
+    "RepositoryChannelDownloadModelCreationParameterDict": "DownloadModelChannelRequestDict",
+}
+
+
 def _generate_data_model_from_json_schema() -> None:
     """Produce Python data model classes from the exported JSON schema file."""
     if not _CACHED_SCHEMA_PATH.exists():
@@ -387,42 +417,73 @@ def _generate_data_model_from_json_schema() -> None:
     model_ast = ast.parse(model_source)
     dict_token_replacements: dict[str, str] = {}
     exported_names: list[str] = []
+    # Scan all nodes in the AST (only in-place node changes are valid here)
+    for node in ast.walk(model_ast):
+        match node:
+            case ast.Name(id=name) as name_node:
+                # Override names when looked up or assigned directly
+                override_name = _DATA_MODEL_NAME_OVERRIDES.get(name, None)
+                if override_name is not None:
+                    name_node.id = override_name
+            case ast.Constant(value=name) as name_constant:
+                # Override names when they appear as type hint forward references
+                override_name = _DATA_MODEL_NAME_OVERRIDES.get(name, None)
+                if override_name is not None:
+                    name_constant.value = override_name
+    # Scan top level nodes only (allows for adding & removing top level nodes)
     for node in model_ast.body:
         match node:
             case ast.ClassDef(name=name):
-                name = node.name
+                # Override names when defining classes
+                override_name = _DATA_MODEL_NAME_OVERRIDES.get(name, None)
+                if override_name is not None:
+                    generated_name = name
+                    name = node.name = override_name
                 exported_names.append(name)
                 if name.endswith("Dict"):
                     struct_name = name.removesuffix("Dict")
                     dict_token_replacements[struct_name] = name
+                    if override_name is not None:
+                        # Fix up docstring reference back to corresponding struct type
+                        expr_node = node.body[0]
+                        assert isinstance(expr_node, ast.Expr)
+                        docstring_node = expr_node.value
+                        assert isinstance(docstring_node, ast.Constant)
+                        docstring = docstring_node.value
+                        assert isinstance(docstring, str)
+                        docstring_node.value = docstring.replace(generated_name, name)
             case ast.Assign(targets=[ast.Name(id=alias)], value=expr):
-                # We don't want to require the specific aliased types for dict inputs
                 match expr:
+                    # For dict fields, replace builtin type aliases with the builtin type names
                     case (
+                        # alias = name
                         ast.Name(id=name)
+                        # alias = Annotated[name, ...]
                         | ast.Subscript(
                             value=ast.Name(id="Annotated"),
                             slice=ast.Tuple(elts=[ast.Name(id=name), *_]),
                         )
                     ):
                         if hasattr(builtins, name):
                             dict_token_replacements[alias] = name
-
+    # Write any AST level changes back to the source file
+    # TODO: Move more changes to the AST rather than relying on raw text replacement
+    _MODEL_PATH.write_text(ast.unparse(model_ast))
     # Additional type union names to be translated
     # Inject the dict versions of required type unions
     # (This is a brute force hack, but it's good enough while there's only a few that matter)
     _single_line_union = (" = ", " | ", "")
-    _multi_line_union = (" = (\n    ", "\n    | ", "\n)")
+    # _multi_line_union = (" = (\n    ", "\n    | ", "\n)")
     _dict_unions = (
         (
-            "ChatMessageData",
+            "AnyChatMessage",
             (
-                "ChatMessageDataAssistant",
-                "ChatMessageDataUser",
-                "ChatMessageDataSystem",
-                "ChatMessageDataTool",
+                "AssistantResponse",
+                "UserMessage",
+                "SystemPrompt",
+                "ToolResultMessage",
             ),
-            _multi_line_union,
+            _single_line_union,
         ),
         (
             "LlmToolUseSetting",