sudosu 0.1.5__py3-none-any.whl

sudosu/core/__init__.py

@@ -0,0 +1,238 @@
+"""Core configuration management for Sudosu."""
+
+import os
+from pathlib import Path
+from typing import Any, Optional
+
+import yaml
+from dotenv import load_dotenv
+
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Default paths - use XDG-compliant location for app data
+# ~/.local/share/sudosu/ for app data (history, etc.)
+# ~/.config/sudosu/ for config (or keep ~/.sudosu/config.yaml for backwards compat)
+APP_DATA_DIR = Path.home() / ".local" / "share" / "sudosu"
+GLOBAL_CONFIG_DIR = Path.home() / ".sudosu"  # Keep for backwards compatibility
+CONFIG_FILE = "config.yaml"
+
+# Default backend URLs (hardcoded, not from env vars)
+DEFAULT_DEV_BACKEND_URL = "ws://localhost:8000/ws"
+DEFAULT_PROD_BACKEND_URL = "wss://sudosu-cli.trysudosu.com/ws"
+
+
+def get_global_config_dir() -> Path:
+    """Get the global configuration directory."""
+    return GLOBAL_CONFIG_DIR
+
+
+def get_app_data_dir() -> Path:
+    """Get the app data directory (for history, cache, etc.)."""
+    return APP_DATA_DIR
+
+
+def get_project_config_dir(cwd: Optional[Path] = None) -> Optional[Path]:
+    """Get the project-specific configuration directory if it exists."""
+    cwd = cwd or Path.cwd()
+    project_config = cwd / ".sudosu"
+    if project_config.exists():
+        return project_config
+    return None
+
+
+def ensure_config_structure() -> Path:
+    """
+    Ensure the global config file exists.
+    
+    NOTE: Only creates config.yaml in ~/.sudosu/, nothing else.
+    All other files (.sudosu/AGENT.md, agents/, etc.) are project-local.
+    """
+    config_dir = get_global_config_dir()
+    
+    # Create config directory (just for config.yaml)
+    config_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Create default config if it doesn't exist
+    config_file = config_dir / CONFIG_FILE
+    if not config_file.exists():
+        default_config = {
+            "mode": "prod",  # default to production
+            "backend_url": DEFAULT_PROD_BACKEND_URL,
+            "dev_backend_url": DEFAULT_DEV_BACKEND_URL,
+            "prod_backend_url": DEFAULT_PROD_BACKEND_URL,
+            "api_key": "",
+            "default_model": "gemini-2.5-pro",
+            "theme": "default",
+        }
+        with open(config_file, "w", encoding="utf-8") as f:
+            yaml.dump(default_config, f, default_flow_style=False)
+    
+    return config_dir
+
+
+def ensure_app_data_dir() -> Path:
+    """Ensure app data directory exists for history, cache, etc."""
+    APP_DATA_DIR.mkdir(parents=True, exist_ok=True)
+    return APP_DATA_DIR
+
+
+def ensure_project_structure(cwd: Optional[Path] = None) -> Path:
+    """
+    Ensure project-local .sudosu/ structure exists with default AGENT.md.
+    
+    This is called when user runs `sudosu` in a directory.
+    Creates:
+      - .sudosu/
+      - .sudosu/AGENT.md (default agent prompt, user-editable)
+      - .sudosu/agents/ (for custom agents)
+    
+    Returns:
+        Path to the project .sudosu directory
+    """
+    from sudosu.core.default_agent import generate_default_agent_md
+    
+    cwd = cwd or Path.cwd()
+    project_config = cwd / ".sudosu"
+    
+    # Create .sudosu/ if it doesn't exist
+    project_config.mkdir(parents=True, exist_ok=True)
+    
+    # Create agents/ subdirectory
+    agents_dir = project_config / "agents"
+    agents_dir.mkdir(exist_ok=True)
+    
+    # Create default AGENT.md if it doesn't exist
+    default_agent_file = project_config / "AGENT.md"
+    if not default_agent_file.exists():
+        default_agent_file.write_text(generate_default_agent_md())
+    
+    return project_config
+
+
+def load_config() -> dict:
+    """Load the global configuration."""
+    config_file = get_global_config_dir() / CONFIG_FILE
+    
+    if not config_file.exists():
+        return {}
+    
+    with open(config_file, "r", encoding="utf-8") as f:
+        return yaml.safe_load(f) or {}
+
+
+def save_config(config: dict) -> None:
+    """Save the global configuration."""
+    config_file = get_global_config_dir() / CONFIG_FILE
+    
+    with open(config_file, "w", encoding="utf-8") as f:
+        yaml.dump(config, f, default_flow_style=False)
+
+
+def get_config_value(key: str, default: Any = None) -> Any:
+    """Get a specific configuration value."""
+    config = load_config()
+    return config.get(key, default)
+
+
+def set_config_value(key: str, value: Any) -> None:
+    """Set a specific configuration value."""
+    config = load_config()
+    config[key] = value
+    save_config(config)
+
+
+def get_mode() -> str:
+    """Get the current mode (dev or prod)."""
+    # Check environment variable first
+    env_mode = os.getenv("SUDOSU_MODE", "").lower()
+    if env_mode in ["dev", "prod"]:
+        return env_mode
+    
+    # Fall back to config file
+    config_mode = get_config_value("mode", "prod")
+    return config_mode if config_mode in ["dev", "prod"] else "prod"
+
+
+def set_mode(mode: str) -> None:
+    """Set the current mode (dev or prod)."""
+    if mode not in ["dev", "prod"]:
+        raise ValueError("Mode must be 'dev' or 'prod'")
+    set_config_value("mode", mode)
+    
+    # Update backend_url based on mode
+    if mode == "dev":
+        url = get_config_value("dev_backend_url", DEFAULT_DEV_BACKEND_URL)
+    else:
+        url = get_config_value("prod_backend_url", DEFAULT_PROD_BACKEND_URL)
+    
+    set_config_value("backend_url", url)
+
+
+def get_backend_url() -> str:
+    """Get the backend WebSocket URL based on current mode."""
+    # Check for direct environment variable override (highest priority)
+    env_url = os.getenv("SUDOSU_BACKEND_URL")
+    if env_url:
+        return env_url
+    
+    # Get current mode
+    mode = get_mode()
+    
+    # Get mode-specific URL
+    if mode == "dev":
+        # Priority: env var > config file > default
+        env_url = os.getenv("SUDOSU_DEV_BACKEND_URL")
+        if env_url:
+            return env_url
+        config_url = get_config_value("dev_backend_url")
+        return config_url if config_url else DEFAULT_DEV_BACKEND_URL
+    else:
+        # Priority: env var > config file > default
+        env_url = os.getenv("SUDOSU_PROD_BACKEND_URL")
+        if env_url:
+            return env_url
+        config_url = get_config_value("prod_backend_url")
+        return config_url if config_url else DEFAULT_PROD_BACKEND_URL
+
+
+def get_agents_dir(project_first: bool = True) -> Optional[Path]:
+    """Get the agents directory (project-specific only).
+    
+    Returns:
+        Path to project's .sudosu/agents/ or None if not in a project
+    """
+    if project_first:
+        project_dir = get_project_config_dir()
+        if project_dir:
+            agents_dir = project_dir / "agents"
+            if agents_dir.exists():
+                return agents_dir
+    
+    return None
+
+
+def get_skills_dir(project_first: bool = True) -> Optional[Path]:
+    """Get the skills directory (project-specific only).
+    
+    Returns:
+        Path to project's .sudosu/skills/ or None if not in a project
+    """
+    if project_first:
+        project_dir = get_project_config_dir()
+        if project_dir:
+            skills_dir = project_dir / "skills"
+            if skills_dir.exists():
+                return skills_dir
+    
+    return None
+
+
+# Export session management
+from sudosu.core.session import (
+    ConversationSession,
+    SessionManager,
+    get_session_manager,
+    reset_session_manager,
+)

sudosu/core/agent_loader.py

@@ -0,0 +1,263 @@
+"""Agent loader - Parse and load AGENT.md files."""
+
+from pathlib import Path
+from typing import Any, Optional
+
+import frontmatter
+
+from sudosu.core.default_agent import CONTEXT_AWARE_PROMPT, SUB_AGENT_CONSULTATION_PROMPT
+
+
+def parse_agent_file(agent_path: Path) -> Optional[dict]:
+    """
+    Parse an AGENT.md file and return the agent configuration.
+    
+    Args:
+        agent_path: Path to the agent directory containing AGENT.md
+    
+    Returns:
+        Agent configuration dict or None if invalid
+    """
+    agent_md = agent_path / "AGENT.md"
+    
+    if not agent_md.exists():
+        return None
+    
+    try:
+        with open(agent_md, "r", encoding="utf-8") as f:
+            post = frontmatter.load(f)
+        
+        # Get description - ensure it's a string
+        description = post.get("description", "")
+        if isinstance(description, list):
+            description = "\n".join(str(item) for item in description)
+        
+        # Get tools list
+        tools = post.get("tools", ["read_file", "write_file", "list_directory"])
+        if isinstance(tools, str):
+            tools = [tools]
+        
+        # Get integrations list
+        integrations = post.get("integrations", [])
+        if isinstance(integrations, str):
+            integrations = [integrations]
+        
+        # Get the base system prompt and append context-aware instructions
+        base_prompt = str(post.content).strip()
+        # Append context awareness to all agents for better conversation handling
+        # Also append consultation instructions so sub-agents know when to consult sudosu
+        system_prompt = base_prompt + CONTEXT_AWARE_PROMPT + SUB_AGENT_CONSULTATION_PROMPT
+        
+        return {
+            "name": str(post.get("name", agent_path.name)),
+            "description": str(description),
+            "model": str(post.get("model", "gemini-2.5-pro")),
+            "tools": tools,
+            "integrations": integrations,
+            "skills": post.get("skills", []),
+            "system_prompt": system_prompt,
+            "path": str(agent_path),
+        }
+    except Exception as e:
+        print(f"Error parsing agent file {agent_md}: {e}")
+        return None
+
+
+def discover_agents(agents_dir: Path) -> list[dict]:
+    """
+    Discover all agents in a directory.
+    
+    Args:
+        agents_dir: Path to the agents directory
+    
+    Returns:
+        List of agent configurations
+    """
+    agents = []
+    
+    if not agents_dir.exists():
+        return agents
+    
+    for item in agents_dir.iterdir():
+        if item.is_dir():
+            agent = parse_agent_file(item)
+            if agent:
+                agents.append(agent)
+    
+    return agents
+
+
+def load_agent_config(agent_name: str, agents_dirs: list[Path]) -> Optional[dict]:
+    """
+    Load a specific agent's configuration by name.
+    
+    Args:
+        agent_name: Name of the agent to load
+        agents_dirs: List of directories to search for agents
+    
+    Returns:
+        Agent configuration dict or None if not found
+    """
+    for agents_dir in agents_dirs:
+        agent_path = agents_dir / agent_name
+        if agent_path.exists():
+            return parse_agent_file(agent_path)
+    
+    return None
+
+
+def create_agent_template(
+    agent_dir: Path,
+    name: str,
+    description: str,
+    system_prompt: str,
+    model: str = "gemini-2.5-pro",
+    tools: Optional[list[str]] = None,
+    integrations: Optional[list[str]] = None,
+) -> Path:
+    """
+    Create a new agent from a template.
+    
+    Args:
+        agent_dir: Directory where the agent will be created
+        name: Agent name
+        description: Agent description
+        system_prompt: Agent system prompt (markdown body)
+        model: Model to use
+        tools: List of allowed tools
+        integrations: List of connected integrations the agent should use
+    
+    Returns:
+        Path to the created agent directory
+    """
+    if tools is None:
+        tools = ["read_file", "write_file", "list_directory"]
+    if integrations is None:
+        integrations = []
+    
+    # Create agent directory
+    agent_path = agent_dir / name
+    agent_path.mkdir(parents=True, exist_ok=True)
+    
+    # Create AGENT.md
+    agent_md = agent_path / "AGENT.md"
+    
+    # Ensure description is properly quoted for YAML (contains colons, etc.)
+    # Use double quotes and escape any internal double quotes
+    safe_description = description.replace('"', '\\"')
+    
+    # Build integrations section if any
+    integrations_yaml = ""
+    if integrations:
+        integrations_yaml = f"""integrations:
+{chr(10).join(f'  - {integration}' for integration in integrations)}"""
+    else:
+        integrations_yaml = "integrations: []"
+    
+    content = f"""---
+name: {name}
+description: "{safe_description}"
+model: {model}
+tools:
+{chr(10).join(f'  - {tool}' for tool in tools)}
+{integrations_yaml}
+skills: []
+---
+
+{system_prompt}
+"""
+    
+    with open(agent_md, "w", encoding="utf-8") as f:
+        f.write(content)
+    
+    return agent_path
+
+
+DEFAULT_WRITER_PROMPT = """# Writer Agent
+
+You are a professional content writer. You help users create well-structured, 
+engaging content including blog posts, articles, and documentation.
+
+## IMPORTANT: Always Save Your Work
+
+**You MUST always save your written content to a file using the write_file tool.**
+- Use kebab-case for filenames (e.g., `my-blog-post.md`)
+- Always use the `.md` extension for blog posts and articles
+- Save the file AFTER you finish writing the content
+
+## Guidelines
+
+1. If the request is clear, start writing immediately
+2. Structure content with clear headings and sections
+3. Use proper markdown formatting
+4. **Always save the final content to a .md file**
+
+## Writing Style
+
+- Clear and concise
+- Engaging but professional
+- Well-researched (when applicable)
+- SEO-friendly for blog posts
+
+## Workflow
+
+1. Write the full content in markdown format
+2. **ALWAYS use write_file to save the content** - this is mandatory
+3. After saving, give a BRIEF confirmation (1-2 sentences max)
+
+## CRITICAL: After Saving a File
+
+**After the write_file tool succeeds, you MUST:**
+- Give a SHORT confirmation like "✓ Saved to filename.md"
+- Optionally add ONE brief insight or suggestion
+- **DO NOT repeat or summarize the content you just saved**
+- **DO NOT show the content again**
+- **Keep your response after saving to 2-3 sentences maximum**
+
+Example good response after saving:
+"✓ Saved your blog post to ai-revolution.md. The article covers 5 key areas where AI is transforming CS. Let me know if you'd like any revisions!"
+
+Example BAD response (DO NOT DO THIS):
+"Here is the blog post about AI... [repeating the entire content]"
+"""
+
+
+DEFAULT_CODER_PROMPT = """# Coder Agent
+
+You are an expert software developer. You help users write, review, and debug code
+across various programming languages and frameworks.
+
+## Guidelines
+
+1. Understand the project context before making changes
+2. Follow existing code style and conventions
+3. Write clean, well-documented code
+4. Consider edge cases and error handling
+
+## Best Practices
+
+- Use meaningful variable and function names
+- Add comments for complex logic
+- Follow SOLID principles
+- Write testable code
+
+## Output Format
+
+When writing code:
+1. Explain the approach first
+2. Show the code with proper formatting
+3. Explain any important decisions
+4. Save files with appropriate names and extensions
+"""
+
+
+AGENT_TEMPLATES = {
+    "writer": {
+        "description": "A helpful writing assistant that creates blog posts, articles, and documentation.",
+        "system_prompt": DEFAULT_WRITER_PROMPT,
+    },
+    "coder": {
+        "description": "An expert developer that helps write, review, and debug code.",
+        "system_prompt": DEFAULT_CODER_PROMPT,
+    },
+}

sudosu/core/connection.py

@@ -0,0 +1,196 @@
+"""WebSocket connection manager for communicating with the backend."""
+
+import asyncio
+import json
+from typing import Any, AsyncGenerator, Callable, Optional
+
+import websockets
+from websockets.client import WebSocketClientProtocol
+
+
+class ConnectionManager:
+    """Manages WebSocket connection to the Sudosu backend."""
+    
+    def __init__(self, backend_url: str):
+        self.backend_url = backend_url
+        self.ws: Optional[WebSocketClientProtocol] = None
+        self._connected = False
+    
+    async def connect(self) -> bool:
+        """Establish connection to the backend.
+        
+        Timeout values increased to support long-running operations
+        like bulk email fetching (100+ emails) and multi-step automations.
+        """
+        try:
+            self.ws = await websockets.connect(
+                self.backend_url,
+                ping_interval=120,   # 2 minutes (was 30s)
+                ping_timeout=60,     # 1 minute (was 10s)  
+                max_size=20_000_000, # 20MB for large responses
+            )
+            self._connected = True
+            return True
+        except Exception as e:
+            self._connected = False
+            raise ConnectionError(f"Failed to connect to backend: {e}")
+    
+    async def disconnect(self) -> None:
+        """Close the connection."""
+        if self.ws:
+            await self.ws.close()
+            self._connected = False
+            self.ws = None
+    
+    @property
+    def is_connected(self) -> bool:
+        """Check if connected to backend."""
+        return self._connected and self.ws is not None
+    
+    async def send(self, message: dict) -> None:
+        """Send a message to the backend."""
+        if not self.is_connected:
+            raise ConnectionError("Not connected to backend")
+        
+        await self.ws.send(json.dumps(message))
+    
+    async def receive(self) -> dict:
+        """Receive a message from the backend."""
+        if not self.is_connected:
+            raise ConnectionError("Not connected to backend")
+        
+        msg = await self.ws.recv()
+        return json.loads(msg)
+    
+    async def invoke_agent(
+        self,
+        agent_config: dict,
+        message: str,
+        cwd: str,
+        session_id: Optional[str] = None,
+        thread_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        on_text: Optional[Callable[[str], None]] = None,
+        on_tool_call: Optional[Callable[[str, dict], Any]] = None,
+        on_status: Optional[Callable[[str], None]] = None,
+        on_special_message: Optional[Callable[[dict], Any]] = None,
+    ) -> AsyncGenerator[dict, None]:
+        """
+        Invoke an agent and stream the response with session context.
+        
+        Args:
+            agent_config: Agent configuration dict
+            message: User message
+            cwd: Current working directory
+            session_id: Session ID for memory continuity across agents
+            thread_id: Thread ID for specific conversation continuity
+            user_id: User ID for integration tools (Gmail, etc.)
+            on_text: Callback for text chunks
+            on_tool_call: Callback for tool calls (should return result)
+            on_status: Callback for status updates
+            on_special_message: Callback for special backend messages (consultation, etc.)
+        
+        Yields:
+            Response messages from the backend
+        """
+        # Build request with session info for memory
+        request = {
+            "type": "invoke",
+            "agent": agent_config,
+            "message": message,
+            "cwd": cwd,
+        }
+        
+        # Include session info if provided
+        if session_id:
+            request["session_id"] = session_id
+        if thread_id:
+            request["thread_id"] = thread_id
+        if user_id:
+            request["user_id"] = user_id
+        
+        # Send invoke request
+        await self.send(request)
+        
+        # Stream responses
+        while True:
+            try:
+                data = await self.receive()
+                msg_type = data.get("type")
+                
+                if msg_type == "text":
+                    if on_text:
+                        on_text(data.get("content", ""))
+                    yield data
+                
+                elif msg_type == "tool_call":
+                    if on_tool_call:
+                        tool_name = data.get("tool")
+                        tool_args = data.get("args", {})
+                        
+                        if on_status:
+                            on_status(f"Executing {tool_name}...")
+                        
+                        # Execute tool and get result
+                        result = await on_tool_call(tool_name, tool_args)
+                        
+                        # Send result back to backend
+                        await self.send({
+                            "type": "tool_result",
+                            "tool": tool_name,
+                            "result": result,
+                        })
+                    yield data
+                
+                elif msg_type == "status":
+                    if on_status:
+                        on_status(data.get("message", ""))
+                    yield data
+                
+                elif msg_type == "done":
+                    yield data
+                    break
+                
+                elif msg_type == "error":
+                    yield data
+                    break
+                
+                elif msg_type == "get_available_agents":
+                    # Backend is requesting available agents for consultation
+                    if on_special_message:
+                        response = await on_special_message(data)
+                        if response:
+                            await self.send(response)
+                    yield data
+                
+                elif msg_type == "consultation_route":
+                    # Backend has decided to route via consultation
+                    if on_special_message:
+                        await on_special_message(data)
+                    yield data
+                
+                elif msg_type == "background_queued":
+                    # Task has been queued for background execution
+                    if on_special_message:
+                        await on_special_message(data)
+                    yield data
+                    # Background tasks complete immediately from client perspective
+                    # The actual work happens in the background
+                    break
+                
+                else:
+                    yield data
+                    
+            except websockets.exceptions.ConnectionClosed:
+                yield {"type": "error", "message": "Connection closed"}
+                break
+            except Exception as e:
+                yield {"type": "error", "message": str(e)}
+                break
+
+
+async def create_connection(backend_url: str) -> ConnectionManager:
+    """Create and connect to the backend."""
+    manager = ConnectionManager(backend_url)
+    await manager.connect()
+    return manager