ask-code/ride_rails.py at main · mogongzi/ask-code

567 lines (475 loc) · 20.9 KB
#!/usr/bin/env python3
Rails Code Analysis CLI using ReAct Agent
Enhanced Rails code analysis tool that uses the ReAct agent architecture
with improved error handling, configuration, and debugging capabilities.
import argparse
import signal
from typing import List, Optional
from rich.console import Console
from rich.markup import escape
from rich.text import Text
# Import core components
from providers import get_provider
from util.command_helpers import handle_special_commands
from util.input_helpers import should_exit_from_input
from chat.session import ChatSession
from llm.clients import StreamingClient, BlockingClient
from llm.types import Provider
# Import agent components
from agent.react_rails_agent import ReactRailsAgent
from agent.config import AgentConfig
from agent.logging import AgentLogger
from agent.reasoning_display import format_complete_reasoning_section
from agent.llm_client import MarkdownStyled
from agent_tool_executor import AgentToolExecutor
from render.explored_display import ExploredDisplay
# Configuration
# Default fallback; overridden per provider below when available
MAX_TOKEN_SIZE = 20000
DEFAULT_URL = "http://127.0.0.1:8000/invoke"
PROMPT_STYLE = "bold green"
console = Console()
def create_streaming_client(
    use_streaming: bool = False, console: Optional[Console] = None, provider_name: str = "bedrock"
    """Create and return streaming or blocking client.
        use_streaming: If True, use StreamingClient (SSE). If False, use BlockingClient (single request).
        console: Rich console for output (used by BlockingClient for spinner)
        provider_name: Provider name string ("bedrock" or "azure")
    Returns:
        StreamingClient or BlockingClient instance
    provider = Provider.from_string(provider_name)
    if use_streaming:
        return StreamingClient(provider=provider)
        return BlockingClient(console=console, provider=provider)
def get_agent_input(
    console, prompt_style, display_string, thinking_mode, user_history, tools_enabled, react_agent=None
    Enhanced input function for the Rails agent.
    Includes better error handling and user feedback.
    from util.simple_pt_input import _create_key_bindings, _prompt_for_input
    def _display_enhanced_instructions(
        token_info: str = None, thinking_mode: bool = False
    ) -> None:
        """Display enhanced usage instructions for the Rails agent."""
        base_instructions = "↵ send    Ctrl+J newline"
        if thinking_mode:
            thinking_part = "/think reasoning [ON]"
        else:
            thinking_part = "/think reasoning"
        instructions = f"{base_instructions}    {thinking_part}    /reasoning    /clear    /status    Esc=cancel"
        safe_token_info = escape(token_info) if token_info else None
        if safe_token_info:
            terminal_width = console.size.width if hasattr(console, "size") else 120
            base_length = len(instructions)
            padding_needed = terminal_width - base_length - len(token_info)
            if padding_needed > 2:
                padding = " " * (padding_needed - 1)
                console.print(f"[dim]{instructions}{padding}{safe_token_info}[/dim]")
            else:
                console.print(f"[dim]{instructions}  {safe_token_info}[/dim]")
        else:
            console.print(f"[dim]{instructions}[/dim]")
    _display_enhanced_instructions(display_string, thinking_mode)
    key_bindings = _create_key_bindings(user_history)
        user_input = _prompt_for_input(key_bindings, user_history, None)
        if not user_input:
            return None, False, thinking_mode, tools_enabled
        # Handle thinking toggle
        if user_input.strip().lower() == "/think":
            thinking_mode = not thinking_mode
            status = "enabled" if thinking_mode else "disabled"
            console.print(f"[yellow]Reasoning mode {status}[/yellow]")
            return None, True, thinking_mode, tools_enabled
        # Handle clear command
        if user_input.strip().lower() == "/clear":
            return "/clear", False, thinking_mode, tools_enabled
        # Handle status command
        if user_input.strip().lower() == "/status":
            return "/status", False, thinking_mode, tools_enabled
        # Tools are always enabled for Rails analysis
        if user_input.strip().lower() == "/tools":
            console.print("[dim]Tools are always enabled for Rails analysis[/dim]")
            return None, False, thinking_mode, tools_enabled
        # Handle reasoning toggle
        if user_input.strip().lower() == "/reasoning":
            if react_agent:
                new_value = not react_agent.config.llm_tracking
                react_agent.config = react_agent.config.update(llm_tracking=new_value)
                status = "enabled" if new_value else "disabled"
                console.print(f"[yellow]Reasoning display {status}[/yellow]")
            else:
                console.print("[dim]Agent not available for reasoning toggle[/dim]")
            return None, True, thinking_mode, tools_enabled
        return user_input, False, thinking_mode, tools_enabled
    except Exception as e:
        console.print(f"[red]Input error: {e}[/red]")
        return None, False, thinking_mode, tools_enabled
    url: str,
    provider,
    project_root: str,
    verbose: bool = False,
    use_streaming: bool = False,
    llm_tracking: bool = False,
    Enhanced interactive Rails code analysis loop with ReAct agent.
        url: LLM endpoint URL
        provider: Provider adapter (bedrock/azure)
        project_root: Rails project root directory
        verbose: Enable verbose mode with detailed logging
        use_streaming: Use streaming API (SSE) vs non-streaming (single request)
        llm_tracking: Show LLM reasoning trail after final answer
    Returns:
        Exit code (0 for success)
    console.rule("Enhanced Rails Analysis Agent")
    # Configure logging - WARNING by default, INFO/DEBUG in verbose mode
    AgentLogger.configure(level="DEBUG" if verbose else "WARNING", console=console)
    # Validate project root
    if not os.path.exists(project_root):
        console.print(
            f"[red]Error: Project directory does not exist: {project_root}[/red]"
        return 1
    # Add usage tracking
    from chat.usage_tracker import UsageTracker
    # Determine provider-specific context length for usage indicator
    provider_context = getattr(provider, "context_length", None)
    usage_max_context = int(provider_context) if provider_context else MAX_TOKEN_SIZE
    usage = UsageTracker(max_tokens_limit=usage_max_context)
    # Extract provider name
    provider_name = (
        provider.__name__.split(".")[-1] if hasattr(provider, "__name__") else "bedrock"
    # Create client (streaming or blocking) with correct provider
    client = create_streaming_client(use_streaming=use_streaming, console=console, provider_name=provider_name)
    if verbose:
        client_type = "streaming (SSE)" if use_streaming else "blocking (single request)"
        console.print(f"[dim]Using {client_type} client (context {usage_max_context} tokens)[/dim]")
    # Create session
    session = ChatSession(
        url=url,
        provider=provider,
        max_tokens=4096,
        timeout=120.0,
        tool_executor=None,
        provider_name=provider_name,
    # Add usage tracker and client to session
    session.usage_tracker = usage
    session.streaming_client = client
    # Initialize ReAct agent
        # Create agent configuration
        # Simplified: trust the LLM to decide when done, only intervene for true infinite loops
        config = AgentConfig(
            project_root=project_root,
            max_react_steps=100,  # Safety limit
            debug_enabled=verbose,
            log_level="DEBUG" if verbose else "WARNING",
            max_exact_repeats=3,  # Only intervene if exact same action repeated 3+ times
            llm_tracking=llm_tracking,
        react_agent = ReactRailsAgent(config=config, session=session, console=console)
        console.print(f"[green]✓ Rails Agent initialized[/green]: {project_root}")
        # Log initialization in verbose mode
        if verbose:
            logger = AgentLogger.get_logger()
            logger.info(f"Created agent for ride_rails.py: {project_root}")
            status = react_agent.get_status()
            logger.debug(f"Agent status: {status}")
        # Show configuration in verbose mode
        if verbose:
            status = react_agent.get_status()
            config = status["config"]
            console.print(
                f"[dim]Config: {config['max_react_steps']} max steps, "
                f"debug={config['debug_enabled']}, "
                f"tools={len(status['tool_registry']['tools_available'])}[/dim]"
    except Exception as e:
        console.print(f"[red]Error: Could not initialize ReAct agent: {e}[/red]")
        if verbose:
            import traceback
            console.print(f"[dim]{traceback.format_exc()}[/dim]")
        return 1
    # Wire tool executor with live Explored display callback
    explored_display = ExploredDisplay(console)
    exploration_tracker = react_agent.state_machine.state.exploration
    def _record_exploration_and_update(tool_name: str, tool_input: dict) -> None:
        """Callback to record exploration and update live display."""
        from agent.exploration_tracker import ExploredType
        # Classify tool and record to tracker
        if tool_name in ["ripgrep", "grep", "search", "enhanced_sql_rails_search"]:
            pattern = tool_input.get("pattern", tool_input.get("sql", tool_input.get("query", "")))
            path = tool_input.get("path", ".")
            exploration_tracker.add_search(pattern, path)
        elif tool_name in ["file_reader", "read_file", "model_analyzer", "controller_analyzer"]:
            file_path = tool_input.get("file_path", tool_input.get("model_name", tool_input.get("controller_name", "")))
            filename = file_path.split("/")[-1] if "/" in str(file_path) else str(file_path)
            exploration_tracker.add_read([filename])
        elif tool_name in ["list_directory", "ls", "directory"]:
            path = tool_input.get("path", ".")
            exploration_tracker.add_list(path)
        else:
            # Default: treat unknown tools as search
            exploration_tracker.add_search(tool_name, ".")
        # Start Live lazily to avoid colliding with other Live spinners (e.g. request spinner).
        if explored_display.live is None:
            explored_display.start_live(exploration_tracker)
        # Update live display with fade-in
        explored_display.add_item_with_fade(exploration_tracker)
        available_tools = react_agent.tool_registry.get_available_tools()
        agent_executor = AgentToolExecutor(available_tools)
        # Recreate client with tool executor, exploration callback, and correct provider
        llm_provider = Provider.from_string(provider_name)
        if use_streaming:
            session.streaming_client = StreamingClient(
                tool_executor=agent_executor,
                on_tool_start=_record_exploration_and_update,
                provider=llm_provider
        else:
            session.streaming_client = BlockingClient(
                tool_executor=agent_executor,
                console=console,
                on_tool_start=_record_exploration_and_update,
                provider=llm_provider
        if verbose:
            console.print(
                f"[dim]Tool executor configured with {len(available_tools)} tools[/dim]"
        if verbose:
            tool_names = list(available_tools.keys())
            console.print(f"[dim]Available tools: {', '.join(tool_names)}[/dim]")
    except Exception as e:
        console.print(
            f"[yellow]Warning: could not attach agent tool executor: {e}[/yellow]"
    # Track UI state
    thinking_mode = False
    tools_enabled = True  # Always enabled for Rails agent
    user_history = []
    while True:
        try:
            # Build enhanced display string
            usage_display = usage.get_display_string()
            mode_indicator = "🧠" if thinking_mode else "🤖"
            # Get user input
            user_input, use_thinking, thinking_mode, tools_enabled = get_agent_input(
                console,
                PROMPT_STYLE,
                usage_display,
                thinking_mode,
                user_history,
                tools_enabled,
                react_agent,
            # Handle exit conditions
            if should_exit_from_input(user_input):
                console.print("[dim]Goodbye! 👋[/dim]")
                return 0
            # Handle clear command
            if user_input and user_input.strip().lower() == "/clear":
                user_history.clear()
                if hasattr(session, "conversation") and session.conversation:
                    session.conversation.clear_history()
                # Clear agent state
                react_agent.state_machine.reset()
                if hasattr(react_agent, "conversation"):
                    react_agent.conversation.clear_history()
                # Clear exploration tracker
                exploration_tracker.clear()
                usage.reset()
                console.print("[green]✓ Conversation and agent state cleared[/green]")
                continue
            # Handle status command
            if user_input and user_input.strip().lower() == "/status":
                status = react_agent.get_status()
                console.print("[bold]Agent Status:[/bold]")
                console.print(f"  Project: {status['config']['project_root']}")
                console.print(
                    f"  Steps completed: {status['state_machine']['current_step']}"
                console.print(
                    f"  Tools used: {len(status['state_machine']['tools_used'])}"
                console.print(
                    f"  Available tools: {len(status['tool_registry']['tools_available'])}"
                console.print(f"  Session queries: {len(user_history)}")
                console.print(f"  Debug mode: {status['config']['debug_enabled']}")
                continue
            # Handle special commands
            if handle_special_commands(
                user_input, None, console, None, None, None, react_agent
                continue
        except (EOFError, KeyboardInterrupt):
            console.print("\n[dim]Goodbye! 👋[/dim]")
            return 0
        except Exception as e:
            if verbose:
                import traceback
                console.print(f"[red]Unexpected error: {e}[/red]")
                console.print(f"[dim]{traceback.format_exc()}[/dim]")
            else:
                console.print(f"[red]Error: {e}[/red]")
            continue
        # Add to history
        if user_input and isinstance(user_input, str):
            user_history.append(user_input)
        # Process through ReAct agent with live Explored display
        try:
            # Clear previous exploration and start fresh
            exploration_tracker.clear()
            exploration_tracker.start()
            # Set thinking mode context if enabled
            if thinking_mode:
                AgentLogger.set_context(thinking_mode=True)
            try:
                response = react_agent.process_message(user_input)
            finally:
                # Mark exploration as complete (changes header to "• Explored")
                exploration_tracker.stop()
                # Update live display to show final "• Explored" state before stopping
                explored_display.update_live(exploration_tracker)
                # Stop live display (content persists on screen in final state)
                explored_display.stop_live()
            # Render final answer after the Explored section.
            if response and response.strip():
                if exploration_tracker.items:
                    console.print()
                console.print(MarkdownStyled(response.strip()))
            # Display ReAct Trace (after answer) if llm_tracking is enabled
            if react_agent.config.llm_tracking:
                cycles = react_agent.get_reasoning_cycles()
                if cycles:
                    format_complete_reasoning_section(cycles, console)
            console.print()  # Add spacing
            # Show usage and session info (only in verbose mode)
            if verbose:
                usage_display = usage.get_display_string()
                session_info = f"Session: {len(user_history)} queries"
                if usage_display:
                    console.print(f"[dim]Usage: {usage_display} • {session_info}[/dim]")
                else:
                    console.print(f"[dim]{session_info}[/dim]")
                # Show analysis steps summary
                    step_summary = react_agent.get_step_summary(limit=8)
                    if step_summary and step_summary.strip() != "No steps recorded.":
                        console.print("[dim]Analysis Steps:[/dim]")
                        for line in step_summary.split("\n"):
                            if line.strip():
                                console.print(f"[dim]  {line}[/dim]")
                    # Show detailed status
                    status = react_agent.get_status()
                    state = status["state_machine"]
                    tools_used = len(state["tools_used"])
                    if tools_used > 0:
                        console.print(
                            f"[dim]Debug: Tools used: {tools_used}, "
                            f"Step: {state['current_step']}, "
                            f"Should stop: {state.get('should_stop', False)}[/dim]"
                except Exception as e:
                    console.print(f"[dim]Debug - Step summary error: {e}[/dim]")
        except Exception as e:
            console.print(f"[red]Agent processing error: {e}[/red]")
            if verbose:
                import traceback
                console.print(f"[dim]{traceback.format_exc()}[/dim]")
def main(argv: Optional[List[str]] = None) -> int:
    """Enhanced CLI entry point with better argument parsing."""
    parser = argparse.ArgumentParser(
        prog="ride-rails",
        description="Enhanced Rails Code Analysis Agent",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
  %(prog)s --project /path/to/rails/app
  %(prog)s --project /path/to/rails/app --verbose
  %(prog)s --project /path/to/rails/app --provider azure --verbose
        """,
    parser.add_argument("--project", required=True, help="Rails project root directory")
    parser.add_argument(
        "--url", default=DEFAULT_URL, help=f"Endpoint URL (default: {DEFAULT_URL})"
    parser.add_argument(
        "--provider",
        default="azure",
        choices=["bedrock", "azure"],
        help="Provider adapter to use (default: azure)",
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable verbose mode with detailed logging (INFO and DEBUG levels)",
    parser.add_argument(
        "--streaming",
        action="store_true",
        help="Use streaming API (SSE) instead of blocking (default: blocking)",
    parser.add_argument(
        "--llm-tracking",
        action="store_true",
        help="Show LLM reasoning trail after the final answer",
    args = parser.parse_args(argv)
    # Setup signal handlers - raise KeyboardInterrupt directly for proper interruption
    def _sigint(_sig, _frm):
        raise KeyboardInterrupt
    def _sigterm(_sig, _frm):
        raise KeyboardInterrupt
    def _sigquit(_sig, _frm):
        raise KeyboardInterrupt
        signal.signal(signal.SIGINT, _sigint)
        signal.signal(signal.SIGTERM, _sigterm)
        signal.signal(signal.SIGQUIT, _sigquit)
    except Exception:
        pass
    # Show startup information
    if args.verbose:
        console.print(f"[dim]Verbose mode enabled[/dim]")
        console.print(f"[dim]Project: {args.project}[/dim]")
        console.print(f"[dim]Provider: {args.provider}[/dim]")
        console.print(f"[dim]Endpoint: {args.url}[/dim]")
    # Get provider and run REPL
        provider = get_provider(args.provider)
        code = repl(
            args.url,
            provider=provider,
            project_root=args.project,
            verbose=args.verbose,
            use_streaming=args.streaming,
            llm_tracking=args.llm_tracking,
        return code
    except Exception as e:
        console.print(f"[red]Startup error: {e}[/red]")
        if args.verbose:
            import traceback
            console.print(f"[dim]{traceback.format_exc()}[/dim]")
        return 1
if __name__ == "__main__":
    raise SystemExit(main())
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

ride_rails.py

Latest commit

History

ride_rails.py

File metadata and controls
