generate_stylus_code
Generate Stylus/Rust smart contract code for Arbitrum using RAG context and version-aware generation. Supports ERC standards and custom contracts with optional tests.
Instructions
Generate Stylus/Rust smart contract code based on requirements. Uses RAG context to provide relevant examples. Supports version-aware generation for different stylus-sdk versions.
Input Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| prompt | Yes | Description of the code to generate | |
| context_query | No | Optional query to retrieve additional context | |
| contract_type | No | Type of contract to generate | |
| include_tests | No | Whether to include unit tests (default: false) | |
| temperature | No | Generation temperature 0-1 (default: 0.2) | |
| target_version | No | Target stylus-sdk version (default: 0.10.0). Use this to generate code for a specific SDK version. | 0.10.0 |
| cargo_toml | No | Optional Cargo.toml content for automatic SDK version detection. If provided, target_version is auto-detected from dependencies. |
Implementation Reference
- The GenerateStylusCodeTool class implements the tool logic for generating and verifying Stylus smart contract code.
class GenerateStylusCodeTool(BaseTool): """ Generates Stylus smart contract code. Uses RAG context to inform code generation with relevant examples. """ MAX_COMPILE_ATTEMPTS = 3 def __init__( self, context_tool: Optional[GetStylusContextTool] = None, compiler_verifier: Optional["CompilerVerifier"] = None, **kwargs, ): """ Initialize the tool. Args: context_tool: GetStylusContextTool for retrieving examples. compiler_verifier: Optional CompilerVerifier for Docker-based cargo check. """ super().__init__(**kwargs) self.context_tool = context_tool or GetStylusContextTool(**kwargs) if compiler_verifier is not None: self.compiler = compiler_verifier elif HAS_COMPILER and CompilerVerifier is not None: self.compiler = CompilerVerifier() else: self.compiler = None def execute( self, prompt: str, context_query: Optional[str] = None, contract_type: Optional[str] = None, include_tests: bool = False, temperature: float = 0.2, target_version: Optional[str] = None, cargo_toml: Optional[str] = None, **kwargs, ) -> dict: """ Generate Stylus smart contract code using template-based generation. Args: prompt: Description of the code to generate. context_query: Optional query to retrieve context. contract_type: Type of contract (token, defi, utility, custom). include_tests: Whether to include unit tests. temperature: Generation temperature (0-1). target_version: Target stylus-sdk version (default: main version). cargo_toml: Optional Cargo.toml content for automatic version detection. Returns: Dict with code, cargo_toml, explanation, dependencies, warnings, context_used, target_version, template_used. """ # Validate input if not prompt or not prompt.strip(): return {"error": "Prompt is required and cannot be empty"} prompt = prompt.strip() warnings = [] # Version detection/selection logic if cargo_toml: detected_version = detect_version_from_cargo_toml(cargo_toml) if detected_version: target_version = detected_version deprecation_warning = get_deprecation_warning(detected_version) if deprecation_warning: warnings.append(deprecation_warning) # Default to main version if not specified if not target_version: target_version = get_main_version() try: # Select appropriate template (version-aware) template = None template_name = "legacy" if HAS_TEMPLATES and select_template: template = select_template( contract_type or "utility", prompt, target_version=target_version ) template_name = template.name # Retrieve relevant context for additional patterns context_used = [] context_text = "" query = context_query or prompt context_result = self.context_tool.execute( query=query, n_results=3, # Reduced since we have a template as base content_type="code", rerank=True, category_boosts=None, # Use default Stylus-focused boosts target_version=target_version, ) if "contexts" in context_result: for ctx in context_result["contexts"]: context_used.append( { "source": ctx["source"], "relevance": ctx["relevance_score"], } ) context_text += ( f"\n--- Example from {ctx['source']} ---\n{ctx['content'][:1500]}\n" ) # Build generation prompt if template: # Use template-based generation user_prompt = self._build_template_prompt( prompt=prompt, template=template, context_text=context_text, include_tests=include_tests, ) system_prompt = get_template_system_prompt(template, target_version) else: # Fallback to legacy generation user_prompt = self._build_prompt( prompt=prompt, contract_type=contract_type, context_text=context_text, include_tests=include_tests, ) system_prompt = get_system_prompt(target_version) messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}, ] response = self._call_llm( messages=messages, temperature=temperature, max_tokens=8192, # Allow longer output for complete contracts ) # Parse response code, cargo_toml_output, explanation = self._parse_template_response( response, template, target_version=target_version ) # Compile-verify-fix loop (if Docker available) compile_verified = False compile_attempts = 0 if self.compiler and self.compiler.is_available() and cargo_toml_output: for attempt in range(self.MAX_COMPILE_ATTEMPTS): compile_attempts = attempt + 1 logger.info(f"Compile check attempt {compile_attempts}") result = self.compiler.verify(code, cargo_toml_output) if result.skipped: logger.info(f"Compile check skipped: {result.skip_reason}") break if result.success: compile_verified = True logger.info("Compile check passed") break # Build fix prompt with structured errors actual_errors = [e for e in result.errors if e.level == "error"] if not actual_errors: compile_verified = True break error_text = format_errors_for_llm(actual_errors, code) guidance = format_fix_guidance(actual_errors) fix_prompt = self._build_fix_prompt(code, error_text, guidance) fix_messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": fix_prompt}, ] fix_response = self._call_llm( messages=fix_messages, temperature=0.1, max_tokens=8192, ) # Parse fixed code fixed_code, _, _ = self._parse_template_response(fix_response, template) if fixed_code and fixed_code != code: code = fixed_code else: warnings.append( f"Compile fix attempt {compile_attempts} did not produce different code" ) break # Extract dependencies with correct versions dependencies = self._extract_dependencies(code, target_version) # Validate code validation_warnings = self._validate_code(code) warnings.extend(validation_warnings) # Derive project name from prompt and fix Cargo.toml/main.rs references # Only include 0.10.0-specific files when targeting 0.10.x+ target_mm = _to_major_minor(target_version) # noqa: F841 main_rs_output = ( (template.main_rs if template else "") if is_at_least_010(target_version) else "" ) stylus_toml_output = ( (template.stylus_toml if template else "") if is_at_least_010(target_version) else "" ) rust_toolchain_toml_output = ( (template.rust_toolchain_toml if template else "") if is_at_least_010(target_version) else "" ) if cargo_toml_output: project_name = self._derive_project_name(prompt) # Fix package name (use underscores for cargo-stylus compatibility) cargo_toml_output = re.sub( r'name\s*=\s*"[^"]+"', f'name = "{project_name}"', cargo_toml_output, ) # Fix main.rs crate reference (print_from_args uses crate name) if main_rs_output: main_rs_output = re.sub( r"(\w+)::print_from_args\b", f"{project_name}::print_from_args", main_rs_output, ) return { "code": code, "cargo_toml": cargo_toml_output, "main_rs": main_rs_output, "stylus_toml": stylus_toml_output, "rust_toolchain_toml": rust_toolchain_toml_output, "explanation": explanation, "dependencies": dependencies, "warnings": warnings if warnings else [], "context_used": context_used, "target_version": target_version, "template_used": template_name, "compile_verified": compile_verified, "compile_attempts": compile_attempts, "disclaimer": TEMPLATE_DISCLAIMER, } except Exception as e: return {"error": f"Code generation failed: {str(e)}"}