aia 1.1.1 → 2.0.0.0.pre.alpha

data/lib/aia/chat_loop.rb

@@ -1,27 +1,53 @@
-# lib/aia/chat_loop.rb
 # frozen_string_literal: true
 
+# lib/aia/chat_loop.rb
+#
+# Thin interactive chat shell for AIA v2.
+# Delegates to StreamingRunner, MentionRouter, and SpecialModeHandler
+# for execution concerns. Owns the REPL loop, directive processing,
+# expert routing, and output coordination.
+
 require "reline"
 require "pm"
 
 module AIA
   class ChatLoop
-    include AIA::SkillUtils
+    include ContentExtractor
 
-    def initialize(chat_processor, ui_presenter, directive_processor)
-      @chat_processor     = chat_processor
-      @ui_presenter       = ui_presenter
+    def initialize(robot, ui_presenter, directive_processor,
+                   session_tracker: nil, alias_registry: nil, filters: {})
+      @robot               = robot
+      @ui_presenter        = ui_presenter
       @directive_processor = directive_processor
+      @tracker             = session_tracker || SessionTracker.new
+      @alias_registry      = alias_registry || ModelAliasRegistry.new
+      @model_switch_handler = ModelSwitchHandler.new(@alias_registry, @ui_presenter)
+      @filters             = filters
+
+      @streaming_runner = StreamingRunner.new
+      @mention_router = MentionRouter.new(
+        ui_presenter: @ui_presenter,
+        tracker: @tracker,
+        streaming_runner: @streaming_runner
+      )
+      @special_mode_handler = SpecialModeHandler.new(
+        robot: @robot,
+        ui_presenter: @ui_presenter,
+        tracker: @tracker
+      )
+      @tool_filter_strategy = ToolFilterStrategy.new(
+        filters: filters,
+        ui_presenter: @ui_presenter
+      )
     end
 
     # Start the interactive chat session
     def start(skip_context_files: false)
       setup_session
-      process_role_context
-      process_skill_context
       process_initial_context(skip_context_files)
-      handle_piped_input
       run_loop
+    rescue StandardError => e
+      AIA.debug_warn("ChatLoop error: #{e.class}: #{e.message}", exc: e)
     ensure
       @ui_presenter.display_chat_end
     end
@@ -29,223 +55,337 @@ module AIA
     private
 
     def setup_session
-      initialize_ui
-      setup_signals
-      Reline::HISTORY.clear
-    end
-
-    def initialize_ui
       puts "\nEntering interactive chat mode..."
       @ui_presenter.display_chat_header
-    end
-
-    def setup_signals
       Signal.trap("INT") { exit }
+      @ui_presenter.load_chat_history
     end
 
-    def process_role_context
-      role = AIA.config.prompts.role
-      return if role.nil? || role.empty?
+    def process_initial_context(skip_context_files)
+      return if skip_context_files || !AIA.config.context_files || AIA.config.context_files.empty?
 
-      prompt_handler = AIA::PromptHandler.new
+      context = AIA.config.context_files.map do |file|
+        File.read(file) rescue "Error reading file: #{file}"
+      end.join("\n\n")
 
-      # fetch_role applies the role file's metadata (including any model: key) to
-      # AIA.config. Preserve the user's chosen model so the chat session keeps
-      # using the model from the CLI, not one embedded in the role file.
-      saved_models = AIA.config.models
-      role_parsed = prompt_handler.fetch_role(role)
-      AIA.config.models = saved_models
+      return if context.empty?
+
+      result, streamed_content, _elapsed = @streaming_runner.run(@robot, context)
+      present_result(result,
+                     streamed_content: streamed_content,
+                     ui_presenter: @ui_presenter)
+    end
 
-      return if role_parsed.nil?
+    def run_loop
+      loop do
+        follow_up_prompt = @ui_presenter.ask_question
+        break if should_exit_chat?(follow_up_prompt)
 
-      role_content = role_parsed.to_s
-      return if role_content.nil? || role_content.strip.empty?
+        log_user_input(follow_up_prompt)
 
-      return unless AIA.client.respond_to?(:chats)
+        processed_prompt = resolve_input(follow_up_prompt)
+        next if processed_prompt.nil?
 
-      system_msg = RubyLLM::Message.new(role: :system, content: role_content)
+        next if model_switch_handled?
+        next if special_routing_handled?(processed_prompt, @robot)
 
-      AIA.client.chats.each_value do |chat|
-        next if chat.messages.any? { |m| m.role == :system }
-        chat.add_message(system_msg)
+        execute_and_present(processed_prompt, @robot)
+        clear_turn_mcp_filter
       end
     end
 
-    def process_skill_context
-      skills = AIA.config.prompts.skills
-      return if skills.nil? || skills.empty?
-
-      skills_dir = AIA.config.skills.dir
-      bodies = Array(skills).filter_map do |skill_name|
-        skill_name = skill_name.to_s.strip
-        next if skill_name.empty?
-
-        skill_path = find_skill_dir(skill_name, skills_dir)
-        next unless skill_path
+    def should_exit_chat?(input)
+      input.nil? || input.strip.downcase == "exit" || input.strip.empty?
+    end
 
-        if File.file?(skill_path)
-          skill_body(File.read(skill_path))
-        else
-          md = File.join(skill_path, 'SKILL.md')
-          skill_body(File.read(md)) if File.exist?(md)
-        end
+    def resolve_input(follow_up_prompt)
+      if follow_up_prompt.strip.start_with?('/')
+        resolve_directive(follow_up_prompt)
+      else
+        parse_prompt_string(follow_up_prompt)
       end
+    end
 
-      return if bodies.empty?
+    def resolve_directive(follow_up_prompt)
+      if @directive_processor.directive?(follow_up_prompt)
+        result = process_directive(follow_up_prompt)
+        # A directive (e.g. /model, /config) may have rebuilt AIA.client;
+        # re-bind so the next prompt uses the new model.
+        rebind_robot_if_rebuilt
+        result
+      else
+        name = follow_up_prompt.strip.split.first
+        @ui_presenter.display_info("Unknown directive: #{name}  (use /help to see available directives)")
+        nil
+      end
+    end
 
-      skill_content = bodies.join("\n\n")
-      response_data = @chat_processor.process_prompt(skill_content)
-      content = response_data.is_a?(Hash) ? response_data[:content] : response_data
-      @chat_processor.output_response(content)
-      @ui_presenter.display_separator
+    def parse_prompt_string(raw)
+      PM.parse_string(raw).to_s
+    rescue StandardError => e
+      @ui_presenter.display_info("Error: #{e.class}: #{e.message}")
+      nil
     end
 
-    def process_initial_context(skip_context_files)
-      return if skip_context_files || !AIA.config.context_files || AIA.config.context_files.empty?
+    def model_switch_handled?
+      return false unless @model_switch_handler.handle(HandlerContext.new(config: AIA.config))
 
-      context = AIA.config.context_files.map do |file|
-        File.read(file) rescue "Error reading file: #{file}"
-      end.join("\n\n")
+      update_robot
+      true
+    end
 
-      return if context.empty?
+    def special_routing_handled?(processed_prompt, active_robot)
+      if @special_mode_handler.handle(processed_prompt)
+        clear_turn_mcp_filter
+        return true
+      end
 
-      response_data = @chat_processor.process_prompt(context)
-      content = response_data.is_a?(Hash) ? response_data[:content] : response_data
+      if @mention_router.handle(HandlerContext.new(robot: active_robot, prompt: processed_prompt))
+        clear_turn_mcp_filter
+        return true
+      end
 
-      @chat_processor.output_response(content)
-      @chat_processor.speak(content)
-      @ui_presenter.display_separator
+      false
     end
 
-    def handle_piped_input
-      return if STDIN.tty?
-      return unless File.exist?("/dev/tty") && File.readable?("/dev/tty") && File.writable?("/dev/tty")
+    def execute_and_present(processed_prompt, active_robot)
+      log_robot_tools(active_robot)
 
-      begin
-        original_stdin = STDIN.dup
-        piped_input = STDIN.read.strip
-        STDIN.reopen("/dev/tty")
-
-        return if piped_input.empty?
+      resolved_tools = @tool_filter_strategy.resolve(processed_prompt)
+      if (AIA.debug? || AIA.verbose?) && resolved_tools
+        puts "\nFiltered tools (#{resolved_tools.size}): #{resolved_tools.join(', ')}"
+      end
 
-        processed_input = PM.parse_string(piped_input).to_s
+      # Plain turn runs the chief directly; aggregation modes run the whole network.
+      turn_target = aggregation_mode? ? active_robot : active_robot.chief
+      result, streamed_content, elapsed = run_streaming_turn(turn_target, processed_prompt, resolved_tools)
+      return unless result
+
+      increment_turn_counter
+      present_result(result,
+                     streamed_content: streamed_content,
+                     prompt: processed_prompt,
+                     elapsed: elapsed,
+                     ui_presenter: @ui_presenter,
+                     tracker: @tracker)
+    end
 
-        response_data = @chat_processor.process_prompt(processed_input)
-        content = response_data.is_a?(Hash) ? response_data[:content] : response_data
+    def run_streaming_turn(turn_target, processed_prompt, resolved_tools)
+      @streaming_runner.run(turn_target, processed_prompt, tools: resolved_tools)
+    rescue StandardError => e
+      @ui_presenter.display_info("Error communicating with AI: #{e.class}: #{e.message}")
+      clear_turn_mcp_filter
+      nil
+    end
 
-        @chat_processor.output_response(content)
-        @chat_processor.speak(content) if AIA.speak?
-        @ui_presenter.display_separator
+    def increment_turn_counter
+      return unless @robot.respond_to?(:memory) && @robot.memory.respond_to?(:data)
 
-        STDIN.reopen(original_stdin)
-      rescue Errno::ENXIO => e
-        warn "Warning: Unable to handle piped input due to TTY unavailability: #{e.message}"
-        return
-      rescue StandardError => e
-        warn "Warning: Error handling piped input: #{e.message}"
-        return
-      end
+      data = @robot.memory.data
+      count = data.respond_to?(:turn_count) ? (data.turn_count || 0) : 0
+      data.turn_count = count + 1
     end
 
-    def run_loop
-      loop do
-        follow_up_prompt = @ui_presenter.ask_question
+    # Clear per-turn MCP server filter so next turn sees all
+    def clear_turn_mcp_filter
+      AIA.turn_state.active_mcp_servers = nil
+    end
 
-        break if follow_up_prompt.nil? || follow_up_prompt.strip.downcase == "exit" || follow_up_prompt.strip.empty?
+    # Update robot reference after a model switch and propagate to sub-components
+    def update_robot
+      @robot = AIA.client
+      @special_mode_handler.robot = @robot
+    end
 
-        if AIA.config.output.file
-          File.open(AIA.config.output.file, "a") do |file|
-            file.puts "\nYou: #{follow_up_prompt}"
-          end
-        end
+    # Re-bind @robot when a directive rebuilt AIA.client into a different robot
+    # (e.g. /model). No-op when the client is unchanged or unset.
+    def rebind_robot_if_rebuilt
+      update_robot if AIA.client && !AIA.client.equal?(@robot)
+    end
 
-        if @directive_processor.directive?(follow_up_prompt)
-          follow_up_prompt = process_directive(follow_up_prompt)
-          next if follow_up_prompt.nil?
-        end
+    # True when the crew should run as an aggregation network (consensus /
+    # parallel / pipeline) rather than routing a plain turn to the chief.
+    def aggregation_mode?
+      cfg = AIA.config
+      cfg.flags.consensus || Array(cfg.pipeline).length > 1 || Array(cfg.models).length > 1
+    end
 
-        begin
-          processed_prompt = PM.parse_string(follow_up_prompt).to_s
-        rescue StandardError => e
-          @ui_presenter.display_info("Error: #{e.class}: #{e.message}")
-          next
-        end
+    def process_directive(follow_up_prompt)
+      directive_output = @directive_processor.process(follow_up_prompt, nil)
 
-        response_data = @chat_processor.process_prompt(processed_prompt)
+      # These directives either mutate state or set a mode flag for the NEXT
+      # prompt.  Display their confirmation and return nil so the loop skips
+      # forwarding them to the robot (which would fire the handler prematurely
+      # on wrapper text and consume the flag before the real prompt arrives).
+      if @directive_processor.state_setting?(follow_up_prompt)
+        @ui_presenter.display_info(directive_output) unless directive_output.nil? || directive_output.strip.empty?
+        return nil
+      end
 
-        if response_data.is_a?(Hash)
-          content = response_data[:content]
-          metrics = response_data[:metrics]
-          multi_metrics = response_data[:multi_metrics]
-        else
-          content = response_data
-          metrics = nil
-          multi_metrics = nil
-        end
+      return nil if directive_output.nil? || directive_output.strip.empty?
 
-        @ui_presenter.display_ai_response(content)
+      puts "\n#{directive_output}\n"
+      "I executed this directive: #{follow_up_prompt}\nHere's the output: #{directive_output}\nLet's continue our conversation."
+    end
 
-        if AIA.config.flags.tokens
-          if multi_metrics
-            @ui_presenter.display_multi_model_metrics(multi_metrics)
-          elsif metrics
-            @ui_presenter.display_token_metrics(metrics)
-          end
-        end
+    # Display token metrics if enabled.
+    # Handles both single-robot results and multi-model network results.
+    #
+    # Token data lives on the raw RubyLLM::Message stored in
+    # RobotResult#raw (RobotLab::Message objects do not carry usage info).
+    def display_metrics(result, elapsed: nil)
+      return unless AIA.config.flags.tokens
 
-        @chat_processor.speak(content)
-        @ui_presenter.display_separator
+      if defined?(SimpleFlow::Result) && result.is_a?(SimpleFlow::Result)
+        display_network_metrics(result)
+        return
       end
-    end
 
-    def process_directive(follow_up_prompt)
-      directive_output = @directive_processor.process(follow_up_prompt, nil)
+      raw = result.respond_to?(:raw) ? result.raw : nil
+      return unless raw && raw.respond_to?(:input_tokens) && raw.input_tokens
+
+      model_id = extract_model_id(raw) || AIA.config.models.first.name
+      metrics = {
+        model_id:      model_id,
+        input_tokens:  raw.input_tokens,
+        output_tokens: raw.output_tokens,
+        elapsed:       elapsed
+      }
+      @ui_presenter.display_token_metrics(metrics)
+    end
 
-      if follow_up_prompt.strip.start_with?("/clear", "/checkpoint", "/restore", "/review", "/context")
-        @ui_presenter.display_info(directive_output) unless directive_output.nil? || directive_output.strip.empty?
-        return nil
+    # Extract per-robot metrics from a network SimpleFlow::Result
+    # and display the multi-model cost table.
+    #
+    # Each robot_result.raw holds the original RubyLLM::Message with
+    # input_tokens, output_tokens, and model_id.
+    # Each robot_result.duration holds the elapsed seconds.
+    # Similarity scores compare each model's response text against the
+    # first model using TF-IDF cosine similarity.
+    def display_network_metrics(flow_result)
+      metrics_list = []
+      response_texts = []
+
+      flow_result.context.each do |task_name, robot_result|
+        next if task_name == :run_params
+        next unless robot_result.respond_to?(:raw)
+
+        raw = robot_result.raw
+        next unless raw && raw.respond_to?(:input_tokens) && raw.input_tokens
+
+        model_id = extract_model_id(raw)
+        display_name = robot_result.respond_to?(:robot_name) ? robot_result.robot_name : task_name.to_s
+        elapsed = robot_result.respond_to?(:duration) ? robot_result.duration : nil
+
+        # Collect response text for similarity scoring
+        text = if robot_result.respond_to?(:reply)
+                 robot_result.reply.to_s
+               elsif robot_result.respond_to?(:content)
+                 robot_result.content.to_s
+               else
+                 ""
+               end
+        response_texts << text
+
+        metrics_list << {
+          model_id:      model_id || display_name,
+          display_name:  display_name,
+          input_tokens:  raw.input_tokens || 0,
+          output_tokens: raw.output_tokens || 0,
+          elapsed:       elapsed
+        }
       end
 
-      return nil if directive_output.nil? || directive_output.strip.empty?
+      return if metrics_list.empty?
 
-      handle_successful_directive(follow_up_prompt, directive_output)
+      # Compute TF-IDF similarity against the first model's response
+      if metrics_list.size > 1
+        scores = SimilarityScorer.score(response_texts)
+        metrics_list.each_with_index { |m, i| m[:similarity] = scores[i] }
+      end
+
+      @ui_presenter.display_multi_model_metrics(metrics_list)
     end
 
-    def handle_successful_directive(follow_up_prompt, directive_output)
-      puts "\n#{directive_output}\n"
-      "I executed this directive: #{follow_up_prompt}\nHere's the output: #{directive_output}\nLet's continue our conversation."
+    # Pull the actual model identifier from a RubyLLM response message
+    # so cost calculation can look up pricing.
+    def extract_model_id(message)
+      return message.model_id if message.respond_to?(:model_id) && message.model_id
+      return message.model    if message.respond_to?(:model)    && message.model
+      nil
     end
 
-    # Parse multi-model response into per-model responses (ADR-002 revised + ADR-005)
-    def parse_multi_model_response(combined_response)
-      return {} if combined_response.nil? || combined_response.empty?
+    # Show the actual tools available to the robot.
+    # Only prints when --debug is enabled.
+    def log_robot_tools(robot)
+      return unless AIA.debug?
 
-      responses = {}
-      current_model = nil
-      current_content = []
+      # active_robot is always a Network; resolve to the chief for tool inspection.
+      target = robot.respond_to?(:chief) ? robot.chief : robot
 
-      combined_response.each_line do |line|
-        if line =~ /^from:\s+(.+)$/
-          if current_model
-            responses[current_model] = current_content.join.strip
-          end
+      local = Array(target.local_tools).map { |t| t.respond_to?(:name) ? t.name : t.class.name }
+      mcp   = Array(target.mcp_tools).map { |t| t.respond_to?(:name) ? t.name : t.class.name }
+
+      $stderr.puts "[DEBUG] Tool filter strategy: #{@tool_filter_strategy.active_strategy_label}"
+      $stderr.puts "[DEBUG] Robot local_tools (#{local.size}): #{local.join(', ')}"
+      $stderr.puts "[DEBUG] Robot mcp_tools (#{mcp.size}): #{mcp.join(', ')}"
+    end
+
+    def log_user_input(input)
+      out_file = AIA.config.output.file
+      return unless out_file
 
-          display_name = $1.strip
-          internal_id = display_name.sub(/\s+\([^)]+\)\s*$/, '')
-          internal_id = internal_id.sub(/\s+#/, '#')
+      File.open(out_file, "a") { |f| f.puts "\nYou: #{input}" }
+    end
 
-          current_model = internal_id
-          current_content = []
-        elsif current_model
-          current_content << line
+    def speak(content)
+      return unless AIA.speak?
+
+      audio   = AIA.config.audio
+      command = audio.speak_command || 'say'
+      env     = {}
+      env['SPEECH_MODEL'] = audio.speech_model if audio.speech_model
+
+      if command == 'say'
+        # Local TTS: say converts and plays in one step
+        run_with_spinner("Speaking...") do
+          if audio.voice && !audio.voice.to_s.strip.empty?
+            system(env, command, '-v', audio.voice, content.to_s)
+          else
+            system(env, command, content.to_s)
+          end
+        end
+      else
+        # Custom TTS script: stage 2 = convert text → audio file,
+        # stage 3 = play the file. AIA passes the output path as $2.
+        require 'tempfile'
+        tmpfile = Tempfile.new(['aia-tts-', '.mp3'])
+        tmpfile.close
+        begin
+          run_with_spinner("Converting to audio...") do
+            system(env, command, content.to_s, tmpfile.path)
+          end
+          if File.size?(tmpfile.path)
+            run_with_spinner("Speaking...") do
+              system('afplay', tmpfile.path)
+            end
+          end
+        ensure
+          tmpfile.unlink
         end
       end
+    rescue StandardError => e
+      $stderr.puts "Warning: Speech failed: #{e.message}"
+    end
 
-      if current_model
-        responses[current_model] = current_content.join.strip
+    def run_with_spinner(message)
+      spinner = TTY::Spinner.new("[:spinner] #{message}", format: :bouncing_ball, output: $stderr)
+      spinner.auto_spin
+      begin
+        yield
+      ensure
+        spinner.stop
       end
-
-      responses
     end
   end
 end