postprocess: Add AssertsTransform and FormattingTransform and GPT tool calling

c2rust-postprocess/postprocess/__init__.py

@@ -21,6 +21,7 @@
 from postprocess.validate import BaselineError, make_validator
 
 DEFAULT_LLM_MODEL = "gemini-3.5-flash"
+DEFAULT_TRANSFORMS = ("comments", "asserts", "formatting")
 
 
 def build_arg_parser() -> argparse.ArgumentParser:
@@ -125,10 +126,11 @@ def build_arg_parser() -> argparse.ArgumentParser:
         type=str,
         required=False,
         action="append",
-        default=["comments"],
+        default=None,
         help=(
             "Transform to apply; pass multiple times to apply multiple transforms "
-            "in sorted order (default: comments)"
+            "in the order provided; duplicate transforms are ignored "
+            f"(default: {', '.join(DEFAULT_TRANSFORMS)})"
         ),
     )
 
@@ -186,13 +188,16 @@ def main(argv: Sequence[str] | None = None):
 
         model = get_model(args.llm_model)
 
-        # sort transform IDs to transforms always run in the same order to
-        # maximize cache hits even if the user passed them in a different order
-        transform_ids = sorted(
-            transform_id.strip()
-            for transform_id in set(args.transform)
-            if transform_id.strip()
+        # De-duplicate transform IDs while preserving their first occurrence.
+        transform_args = args.transform or DEFAULT_TRANSFORMS
+        transform_ids = list(
+            dict.fromkeys(
+                transform_id.strip()
+                for transform_id in transform_args
+                if transform_id.strip()
+            )
         )
+
         transforms = [
             get_transform_by_id(
                 transform_id,

c2rust-postprocess/postprocess/models/gpt.py

@@ -1,11 +1,25 @@
+import inspect
+import json
 from collections.abc import Callable, Iterable
-from typing import Any
+from typing import Any, Protocol, cast
 
 from openai import OpenAI
+from openai.types.responses import (
+    FunctionToolParam,
+    ResponseFunctionToolCall,
+    ResponseInputParam,
+)
+from openai.types.responses.response_input_param import FunctionCallOutput
 
 from postprocess.models import AbstractGenerativeModel
 
 
+class NamedCallable(Protocol):
+    __name__: str
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
 class GPTModel(AbstractGenerativeModel):
     def __init__(
         self,
@@ -22,13 +36,100 @@ def generate_with_tools(
         tools: Iterable[Callable[..., Any]] = (),
         max_tool_loops: int = 5,
     ) -> str:
-        # TODO: implement tool calling support
-        assert not tools, "Tool calling not yet implemented for GPTModel"
+        tools = [self._named_tool(tool) for tool in tools]
+        tool_schemas = [self._tool_schema(tool) for tool in tools]
+        tool_by_name = {tool.__name__: tool for tool in tools}
+
+        if tool_schemas:
+            response = self.client.responses.create(
+                model=self.id,
+                input=messages[0]["content"],
+                max_tool_calls=max_tool_loops,
+                tools=tool_schemas,
+            )
+        else:
+            response = self.client.responses.create(
+                model=self.id,
+                input=messages[0]["content"],
+                max_tool_calls=max_tool_loops,
+            )
 
-        response = self.client.responses.create(
-            model=self.id,
-            input=messages[0]["content"],
-            max_tool_calls=max_tool_loops,
-        )
+        for _ in range(max_tool_loops):
+            tool_calls = [
+                cast(ResponseFunctionToolCall, item)
+                for item in response.output
+                if getattr(item, "type", None) == "function_call"
+            ]
+            if not tool_calls:
+                return response.output_text
+
+            tool_outputs: ResponseInputParam = [
+                FunctionCallOutput(
+                    type="function_call_output",
+                    call_id=tool_call.call_id,
+                    output=self._call_tool(tool_call, tool_by_name),
+                )
+                for tool_call in tool_calls
+            ]
+            response = self.client.responses.create(
+                model=self.id,
+                input=tool_outputs,
+                previous_response_id=response.id,
+                max_tool_calls=max_tool_loops,
+                tools=tool_schemas,
+            )
 
         return response.output_text
+
+    def _named_tool(self, tool: Callable[..., Any]) -> NamedCallable:
+        if not hasattr(tool, "__name__"):
+            raise TypeError(f"Tool must be a named function: {tool!r}")
+        return cast(NamedCallable, tool)
+
+    def _tool_schema(self, tool: NamedCallable) -> FunctionToolParam:
+        signature = inspect.signature(tool)
+        properties: dict[str, object] = {}
+        required: list[str] = []
+        for name, parameter in signature.parameters.items():
+            properties[name] = {
+                "type": self._json_schema_type(parameter.annotation),
+            }
+            if parameter.default is inspect.Parameter.empty:
+                required.append(name)
+
+        return {
+            "type": "function",
+            "name": tool.__name__,
+            "description": inspect.getdoc(tool) or f"Call `{tool.__name__}`.",
+            "parameters": {
+                "type": "object",
+                "properties": properties,
+                "required": required,
+                "additionalProperties": False,
+            },
+            "strict": False,
+        }
+
+    def _json_schema_type(self, annotation: Any) -> str:
+        if annotation is bool:
+            return "boolean"
+        if annotation is int:
+            return "integer"
+        if annotation is float:
+            return "number"
+        return "string"
+
+    def _call_tool(
+        self,
+        tool_call: ResponseFunctionToolCall,
+        tool_by_name: dict[str, NamedCallable],
+    ) -> str:
+        if tool_call.name not in tool_by_name:
+            raise ValueError(f"Unknown tool call: {tool_call.name}")
+
+        arguments = json.loads(tool_call.arguments or "{}")
+        if not isinstance(arguments, dict):
+            raise ValueError(f"Tool call arguments must be an object: {arguments}")
+
+        result = tool_by_name[tool_call.name](**arguments)
+        return result if isinstance(result, str) else json.dumps(result)

c2rust-postprocess/postprocess/transforms/__init__.py

@@ -1,7 +1,9 @@
 from postprocess.cache import AbstractCache
 from postprocess.models import AbstractGenerativeModel
+from postprocess.transforms.asserts import AssertsTransform
 from postprocess.transforms.base import AbstractTransform
 from postprocess.transforms.comments import CommentsTransform
+from postprocess.transforms.formatting import FormattingTransform
 
 
 def get_transform_by_id(
@@ -13,5 +15,9 @@ def get_transform_by_id(
     match id.lower():
         case "comments":
             return CommentsTransform(cache=cache, model=model)
+        case "asserts":
+            return AssertsTransform(cache=cache, model=model)
+        case "formatting":
+            return FormattingTransform(cache=cache, model=model)
         case _:
             raise ValueError(f"Unsupported transform: {id}")

c2rust-postprocess/postprocess/transforms/asserts.py

@@ -0,0 +1,133 @@
+import logging
+from collections.abc import Callable
+from pathlib import Path
+from textwrap import dedent
+
+from postprocess.cache import AbstractCache
+from postprocess.definitions import CDefinition
+from postprocess.models import AbstractGenerativeModel
+from postprocess.transforms.base import AbstractTransform, TransformError
+from postprocess.utils import remove_backticks
+
+SYSTEM_INSTRUCTION = (
+    "You are a helpful assistant that rewrites c2rust-transpiled assert patterns "
+    "into idiomatic Rust assert! macros."
+)
+
+
+class AssertsTransformPrompt:
+    c_function: str
+    rust_function: str
+    prompt_text: str
+    identifier: str
+
+    __slots__ = ("c_function", "rust_function", "prompt_text", "identifier")
+
+    def __init__(
+        self, c_function: str, rust_function: str, prompt_text: str, identifier: str
+    ):
+        self.c_function = c_function
+        self.rust_function = rust_function
+        self.prompt_text = prompt_text
+        self.identifier = identifier
+
+    def __str__(self) -> str:
+        return (
+            self.prompt_text
+            + "\n\n"
+            + "C function:\n```c\n"
+            + self.c_function
+            + "\n```\n\n"
+            + "Rust function:\n```rust\n"
+            + self.rust_function
+            + "\n```\n"
+        )
+
+
+class AssertsTransform(AbstractTransform):
+    def __init__(self, cache: AbstractCache, model: AbstractGenerativeModel):
+        super().__init__(SYSTEM_INSTRUCTION, cache, model)
+
+    @staticmethod
+    def get_validation_fn(expected_assert_count: int) -> Callable[[str], str]:
+        def validate_response(rust_fn: str) -> str:
+            rust_fn = remove_backticks(rust_fn)
+
+            if "__assert_fail(" in rust_fn:
+                return (
+                    "FAILURE: Rust function still contains __assert_fail. "
+                    "Rewrite those into assert! calls. "
+                    "Reply with the full Rust function definition only; "
+                    "say nothing else."
+                )
+
+            actual_assert_count = rust_fn.count("assert!(")
+            if actual_assert_count < expected_assert_count:
+                return (
+                    "FAILURE: Missing rewritten assert! calls. "
+                    f"Expected at least {expected_assert_count}, "
+                    f"got {actual_assert_count}. "
+                    "Reply with the full Rust function definition only; "
+                    "say nothing else."
+                )
+
+            return "SUCCESS: Asserts transformed correctly!"
+
+        return validate_response
+
+    def try_apply_ident(
+        self,
+        rust_source_file: Path,
+        rust_definition: str,
+        c_definition: CDefinition,
+        identifier: str,
+    ) -> str | None:
+        _ = rust_source_file
+        expected_assert_count = rust_definition.count("__assert_fail(")
+        if expected_assert_count == 0:
+            logging.info(
+                f"{self.__class__.__name__}: "
+                f"Skipping function without transpiled asserts: {identifier}"
+            )
+            return
+
+        prompt_text = """
+        Rewrite the Rust function below by replacing transpiled C assert-macro
+        expansions (which call __assert_fail) with idiomatic Rust assert! calls.
+
+        Requirements:
+        - Preserve function behavior.
+        - Preserve formatting and indentation.
+        - Keep all non-assert logic unchanged.
+        - Return the full Rust function definition only; say nothing else.
+        """
+        prompt_text = dedent(prompt_text).strip()
+
+        prompt = AssertsTransformPrompt(
+            c_function=c_definition.effective,
+            rust_function=rust_definition,
+            prompt_text=prompt_text,
+            identifier=identifier,
+        )
+
+        messages = [{"role": "user", "content": str(prompt)}]
+        validate_response = self.get_validation_fn(expected_assert_count)
+
+        def check(response: str) -> str:
+            if response.strip() == "":
+                raise TransformError("model returned an empty response")
+
+            validation_result = validate_response(response)
+            if not validation_result.startswith("SUCCESS"):
+                raise TransformError(
+                    f"model response for {identifier} failed validation: "
+                    f"{validation_result}\nResponse was:\n{response}"
+                )
+            return remove_backticks(response)
+
+        return self.generate(
+            identifier,
+            messages,
+            check,
+            tools=[validate_response],
+        )

c2rust-postprocess/postprocess/transforms/base.py

@@ -1,6 +1,6 @@
 import logging
 import re
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 from dataclasses import dataclass, field
 from functools import partial
 from pathlib import Path
@@ -119,6 +119,7 @@ def generate(
         identifier: str,
         messages: list[dict[str, Any]],
         check: Callable[[str], str],
+        tools: Iterable[Callable[..., Any]] = (),
     ) -> str | None:
         """
         Model call with caching, validation, and retries. `check` returns the
@@ -154,7 +155,7 @@ def generate(
 
         for attempt in range(self.max_attempts):
             try:
-                response = self.model.generate_with_tools(messages)
+                response = self.model.generate_with_tools(messages, tools=tools)
                 if response is None:
                     raise TransformError(f"model returned no response for {identifier}")
                 result = check(response)

c2rust-postprocess/postprocess/transforms/comments.py

@@ -40,10 +40,10 @@ def __str__(self) -> str:
             + "\n```\n\n"
             + "C function:\n```c\n"
             + self.c_function
-            + "```\n\n"
+            + "\n```\n\n"
             + "Rust function:\n```rust\n"
             + self.rust_function
-            + "```\n"
+            + "\n```\n"
         )