wehub-resource-sync
143 lines
4.9 KiB
Python
143 lines
4.9 KiB
Python
"""
|
|
WriteCode tool - Replace model.py in one operation.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import aiofiles
|
|
|
|
from agent.tools.base import (
|
|
BaseDeclarativeTool,
|
|
BoundFileToolInvocation,
|
|
ToolParamsModel,
|
|
ToolResult,
|
|
make_tool_schema,
|
|
)
|
|
from agent.tools.model_contract import missing_required_model_contract
|
|
|
|
|
|
class WriteCodeParams(ToolParamsModel):
|
|
"""Parameters for write_code tool"""
|
|
|
|
code: str
|
|
|
|
|
|
class WriteCodeInvocation(BoundFileToolInvocation[WriteCodeParams, str]):
|
|
"""Invocation for replacing the full model file."""
|
|
|
|
def get_description(self) -> str:
|
|
preview = self.params.code[:50].replace("\n", "\\n")
|
|
if len(self.params.code) > 50:
|
|
preview += "..."
|
|
return f"Rewrite current model.py file: '{preview}'"
|
|
|
|
async def execute(self) -> ToolResult:
|
|
try:
|
|
if not self.file_path:
|
|
return ToolResult(error="file_path is required")
|
|
|
|
missing = missing_required_model_contract(self.params.code)
|
|
if missing:
|
|
return ToolResult(
|
|
error=(
|
|
"write_file must include the complete top-level model.py contract: "
|
|
+ ", ".join(missing)
|
|
)
|
|
)
|
|
|
|
validation = self._validate_python_syntax(
|
|
self.params.code, self.file_path or "<string>"
|
|
)
|
|
|
|
async with aiofiles.open(self.file_path, mode="w", encoding="utf-8") as f:
|
|
await f.write(self.params.code)
|
|
|
|
return ToolResult(output="model.py rewritten successfully", compilation=validation)
|
|
except FileNotFoundError:
|
|
return ToolResult(error=f"File {self.file_path} not found")
|
|
except Exception as exc:
|
|
return ToolResult(error=f"Error writing code: {str(exc)}")
|
|
|
|
def _validate_python_syntax(self, full_code: str, filename: str) -> dict:
|
|
try:
|
|
compile(full_code, filename, "exec")
|
|
return {
|
|
"status": "success",
|
|
"error": None,
|
|
}
|
|
except SyntaxError as exc:
|
|
error_msg = f"Syntax error: {exc.msg} (line {exc.lineno})"
|
|
return {
|
|
"status": "error",
|
|
"error": error_msg,
|
|
"error_line": exc.lineno,
|
|
}
|
|
except Exception as exc:
|
|
return {
|
|
"status": "error",
|
|
"error": f"Validation error: {str(exc)}",
|
|
}
|
|
|
|
|
|
class WriteFileParams(ToolParamsModel):
|
|
"""Parameters for Gemini-style write_file alias."""
|
|
|
|
content: str
|
|
path: str | None = None
|
|
|
|
|
|
class WriteFileInvocation(BoundFileToolInvocation[WriteFileParams, str]):
|
|
"""Invocation for full model.py rewrites under the official-style name."""
|
|
|
|
def get_description(self) -> str:
|
|
preview = self.params.content[:50].replace("\n", "\\n")
|
|
if len(self.params.content) > 50:
|
|
preview += "..."
|
|
return f"Write current model.py file: '{preview}'"
|
|
|
|
async def execute(self) -> ToolResult:
|
|
mapped = WriteCodeParams(code=self.params.content)
|
|
invocation = WriteCodeInvocation(mapped)
|
|
invocation.bind_file_path(self.file_path or "")
|
|
return await invocation.execute()
|
|
|
|
|
|
class WriteFileTool(BaseDeclarativeTool):
|
|
"""Gemini-style write_file tool scoped to the full model file."""
|
|
|
|
def __init__(self) -> None:
|
|
schema = make_tool_schema(
|
|
name="write_file",
|
|
description=(
|
|
"Replace the entire current bound `model.py` artifact.\n\n"
|
|
"This is the Gemini-style full rewrite tool. It writes the full visible model script, "
|
|
"including imports and `object_model = build_object_model()`.\n\n"
|
|
"Use this when a targeted `replace` would be awkward or when you intend to rewrite the "
|
|
"whole model file from scratch."
|
|
),
|
|
parameters={
|
|
"content": {
|
|
"type": "string",
|
|
"description": (
|
|
"Full replacement content for model.py. Include imports, top-level "
|
|
"build_object_model(), run_tests(), and object_model = build_object_model()."
|
|
),
|
|
},
|
|
"path": {
|
|
"type": "string",
|
|
"description": (
|
|
"Optional virtual path for parity with Gemini CLI. "
|
|
'If provided, use `"model.py"`.'
|
|
),
|
|
},
|
|
},
|
|
required=["content"],
|
|
)
|
|
super().__init__("write_file", schema)
|
|
|
|
async def build(self, params: dict) -> WriteFileInvocation:
|
|
validated = WriteFileParams(**params)
|
|
if validated.path not in {None, "model.py"}:
|
|
raise ValueError("write_file only supports path='model.py' in this harness")
|
|
return WriteFileInvocation(validated)
|