refactor docstring and remove extra code

backend/core/prompts/__init__.py

@@ -1,11 +1,16 @@
 """
-Suna.so Modular Prompt System
+Modular prompt system for Suna.so AI agents.
 
-A structured, maintainable system for building AI agent prompts using
+Provides prompt generation using YAML/JSON components with 44.6% token reduction.
-YAML configuration and JSON schemas.
 """
 
-from .assembler import PromptAssembler, get_system_prompt
+from .assembler import PromptAssembler
+from .prompt import get_system_prompt, get_custom_prompt, SYSTEM_PROMPT
 
 __version__ = "2.0.0"
-__all__ = ["PromptAssembler", "get_system_prompt"]
+__all__ = [
+    "PromptAssembler",
+    "get_system_prompt",
+    "get_custom_prompt",
+    "SYSTEM_PROMPT",
+]

backend/core/prompts/assembler.py

@@ -1,33 +1,21 @@
 """
-Modular System Prompt Assembler
+Prompt assembler for building system prompts from modular components.
 
-Dynamically assembles prompts from YAML/JSON components for maximum efficiency.
+Loads and combines YAML configurations, JSON tool schemas, and instruction
-This module provides the core functionality for building structured system prompts
+templates into complete system prompts with caching for performance.
-from modular configuration files, tool schemas, and context templates.
 """
 
 import json
-import os
 from pathlib import Path
 from typing import Dict, List, Optional, Any
 import yaml
-from functools import lru_cache
 
 
 class PromptAssembler:
-    """
+    """Assembles system prompts from YAML/JSON components with caching."""
-    Dynamic prompt assembler that builds system prompts from modular YAML/JSON components.
-    Implements caching, validation, and conditional loading for token efficiency.
-
-    Attributes:
-        base_path: Root directory of the prompt system
-        config_path: Path to configuration files (YAML)
-        schemas_path: Path to tool schema definitions (JSON)
-        templates_path: Path to context templates (YAML)
-    """
 
     def __init__(self, base_path: Optional[str] = None):
-        """Initialize the prompt assembler with base path to prompt_system directory."""
+        """Initialize assembler with base path to prompt directory."""
         if base_path is None:
             base_path = Path(__file__).parent
         self.base_path = Path(base_path)
@@ -35,12 +23,11 @@ class PromptAssembler:
         self.schemas_path = self.base_path / "schemas"
         self.templates_path = self.base_path / "templates"
 
-        # Cache for loaded components
         self._cache = {}
-        self._assembly_cache = {}  # Cache for assembled prompts
+        self._assembly_cache = {}
 
     def _load_yaml(self, file_path: Path) -> Dict[str, Any]:
-        """Load and parse YAML file with caching."""
+        """Load and cache YAML file."""
         cache_key = str(file_path)
         if cache_key in self._cache:
             return self._cache[cache_key]
@@ -52,7 +39,7 @@ class PromptAssembler:
         return data
 
     def _load_json(self, file_path: Path) -> Dict[str, Any]:
-        """Load and parse JSON file with caching."""
+        """Load and cache JSON file."""
         cache_key = str(file_path)
         if cache_key in self._cache:
             return self._cache[cache_key]
@@ -64,10 +51,9 @@ class PromptAssembler:
         return data
 
     def load_config(self) -> Dict[str, Any]:
-        """Load main system configuration."""
+        """Load main system configuration with includes."""
         main_config = self._load_yaml(self.config_path / "system.yaml")
 
-        # Load included configs
         config = {}
         config.update(main_config)
 
@@ -80,37 +66,24 @@ class PromptAssembler:
         return config
 
     def load_tool_schema(self, schema_name: str) -> Dict[str, Any]:
-        """
+        """Load tool schema by name."""
-        Load a specific tool schema by name.
-
-        Args:
-            schema_name: Name of the schema file (without .json extension)
-                        e.g., 'files', 'web', 'design', 'agents', 'knowledge_base'
-
-        Returns:
-            Dictionary containing the JSON schema definition
-        """
         schema_file = self.schemas_path / f"{schema_name}.json"
         return self._load_json(schema_file)
 
     def load_template(self, template_name: str) -> Dict[str, Any]:
-        """Load a specific template by name."""
+        """Load template by name."""
         template_file = self.templates_path / f"{template_name}.yaml"
         return self._load_yaml(template_file)
 
     def _format_agent_identity(self, config: Dict[str, Any]) -> str:
         """Format agent identity section."""
         agent = config.get("agent", {})
-        behavior = config.get("behavior", {})
         principles = config.get("principles", [])
 
         sections = []
-
-        # Agent identity
         sections.append(f"# Agent Identity")
         sections.append(f"{agent.get('identity_statement', '')}\n")
 
-        # Core principles
         if principles:
             sections.append("# Core Principles")
             for principle in principles:
@@ -128,19 +101,16 @@ class PromptAssembler:
         sections = []
         sections.append("# Execution Environment")
 
-        # Base environment
         if "base" in env:
             base = env["base"]
             sections.append(f"- OS: {base.get('os', 'N/A')}")
             sections.append(f"- Python: {base.get('python_version', 'N/A')}")
 
-        # Capabilities summary
         sections.append("\n# Operational Capabilities")
         for cap_name, cap_value in caps.items():
             if isinstance(cap_value, list):
                 formatted = ', '.join(str(item) for item in cap_value)
             elif isinstance(cap_value, dict):
-                # Format dict as key: value pairs
                 formatted = ', '.join(f"{k}: {v}" for k, v in cap_value.items())
             else:
                 formatted = str(cap_value)
@@ -149,7 +119,7 @@ class PromptAssembler:
         return "\n".join(sections)
 
     def _format_tool_schema(self, schema: Dict[str, Any]) -> str:
-        """Format tool schema into concise documentation."""
+        """Format tool schema into documentation."""
         tools = schema.get("tools", [])
 
         sections = []
@@ -161,7 +131,6 @@ class PromptAssembler:
             sections.append(f"\n### {name}")
             sections.append(f"{desc}")
 
-            # Parameters
             params = tool.get("parameters", {})
             if "properties" in params:
                 sections.append("\n**Parameters:**")
@@ -171,7 +140,6 @@ class PromptAssembler:
                     param_desc = param_info.get("description", "")
                     sections.append(f"- `{param_name}`: {param_info.get('type', 'any')}{req} - {param_desc}")
 
-            # Critical notes
             if "critical_notes" in tool:
                 sections.append("\n**Critical Notes:**")
                 for note in tool["critical_notes"]:
@@ -180,10 +148,9 @@ class PromptAssembler:
         return "\n".join(sections)
 
     def _format_template(self, template: Dict[str, Any]) -> str:
-        """Format template into prompt instructions."""
+        """Format template into instructions."""
         sections = []
 
-        # Critical rules
         if "critical_rules" in template:
             sections.append("## Critical Rules")
             for rule in template["critical_rules"]:
@@ -191,7 +158,6 @@ class PromptAssembler:
                 if "reason" in rule:
                     sections.append(f"  Reason: {rule['reason']}")
 
-        # Example workflows
         if "example_workflows" in template:
             sections.append("\n## Workflows")
             for workflow_name, steps in template["example_workflows"].items():
@@ -207,35 +173,29 @@ class PromptAssembler:
         include_templates: Optional[List[str]] = None
     ) -> str:
         """
-        Assemble complete system prompt based on requirements.
+        Assemble system prompt from specified tools and templates.
 
         Args:
-            include_tools: List of tool schemas to include (e.g., ['file_operations', 'web_operations'])
+            include_tools: Tool schemas to include
-            include_templates: List of templates to include (e.g., ['file_ops', 'browser'])
+            include_templates: Templates to include
 
         Returns:
             Assembled system prompt string
         """
-        # Create cache key
         cache_key = (
             tuple(include_tools) if include_tools else None,
             tuple(include_templates) if include_templates else None
         )
 
-        # Check cache
         if cache_key in self._assembly_cache:
             return self._assembly_cache[cache_key]
 
         sections = []
 
-        # Load configuration
         config = self.load_config()
-
-        # Format core sections
         sections.append(self._format_agent_identity(config))
         sections.append(self._format_environment(config))
 
-        # Load and format tool schemas
         if include_tools:
             sections.append("\n# Available Tools")
             for tool_name in include_tools:
@@ -245,7 +205,6 @@ class PromptAssembler:
                 except FileNotFoundError:
                     print(f"Warning: Tool schema '{tool_name}' not found")
 
-        # Load and format templates
         if include_templates:
             sections.append("\n# Specialized Instructions")
             for template_name in include_templates:
@@ -255,24 +214,18 @@ class PromptAssembler:
                 except FileNotFoundError:
                     print(f"Warning: Template '{template_name}' not found")
 
-        # Assemble and cache
         result = "\n\n".join(sections)
         self._assembly_cache[cache_key] = result
         return result
 
     def get_full_prompt(self) -> str:
-        """
+        """Get complete system prompt with all tools and templates."""
-        Get the complete system prompt with all tools and templates.
-
-        Returns:
-            Complete assembled system prompt string
-        """
         return self.assemble_prompt(
             include_tools=["files", "knowledge_base", "web", "agents", "design"],
             include_templates=["files", "web", "browser", "design", "agents"]
         )
 
     def clear_cache(self):
-        """Clear the internal cache."""
+        """Clear internal caches."""
         self._cache.clear()
         self._assembly_cache.clear()

backend/core/prompts/prompt.py

@@ -1,44 +1,18 @@
 """
-System Prompt Module
+System prompt module for Suna.so AI agents.
 
-This module provides system prompts for Suna.so AI agents using a modular
+Provides functions to generate complete or custom system prompts using a
-YAML + JSON approach for efficient prompt generation.
+modular YAML + JSON architecture. Reduces token usage by 44.6% compared
-
+to the original monolithic prompt.
-The modular system:
-- Reduces token usage by 44.6% (3,896 → 2,158 tokens)
-- Allows custom loading of specific tools/templates
-- Provides better maintainability and extensibility
-
-Usage:
-    from backend.core.prompts.prompt import get_system_prompt
-
-    # Get complete system prompt
-    prompt = get_system_prompt()
-
-    # Get custom prompt with specific tools
-    from backend.core.prompts.prompt import get_custom_prompt
-    prompt = get_custom_prompt(
-        include_tools=['files', 'web'],
-        include_templates=['files', 'web']
-    )
 """
 
 from .assembler import PromptAssembler
 
-# Create assembler instance
 _assembler = PromptAssembler()
 
 
 def get_system_prompt() -> str:
-    """
+    """Get the complete system prompt with all capabilities."""
-    Get the complete system prompt for the AI agent.
-
-    Returns:
-        Assembled system prompt string with all capabilities
-
-    Examples:
-        >>> prompt = get_system_prompt()
-    """
     return _assembler.get_full_prompt()
 
 
@@ -50,20 +24,11 @@ def get_custom_prompt(
     Get a custom system prompt with specific tools and templates.
 
     Args:
-        include_tools: List of tool schemas to include.
+        include_tools: Tool schemas to include
-            Available: ['files', 'knowledge_base', 'web', 'agents', 'design']
+        include_templates: Templates to include
-        include_templates: List of templates to include.
-            Available: ['files', 'web', 'browser', 'design', 'agents']
 
     Returns:
-        Custom assembled system prompt string
+        Assembled system prompt string
-
-    Examples:
-        >>> # Only file operations and web tools
-        >>> prompt = get_custom_prompt(
-        ...     include_tools=['files', 'web'],
-        ...     include_templates=['files', 'web']
-        ... )
     """
     return _assembler.assemble_prompt(
         include_tools=include_tools,
@@ -71,13 +36,4 @@ def get_custom_prompt(
     )
 
 
-# Default system prompt constant
+SYSTEM_PROMPT = get_system_prompt()
-SYSTEM_PROMPT = get_system_prompt()
-
-
-# Export public API
-__all__ = [
-    'get_system_prompt',
-    'get_custom_prompt',
-    'SYSTEM_PROMPT',
-]