aia 1.1.1 → 2.0.0.0.pre.alpha

data/lib/aia/config.rb

@@ -18,7 +18,7 @@
 # 5. Embedded directives (/config)
 #
 # When -c / --config-file is used, it REPLACES sources 2 and 3:
-#   defaults → config file → CLI arguments → embedded directives
+#   defaults -> config file -> CLI arguments -> embedded directives
 
 require 'myway_config'
 require 'yaml'
@@ -28,7 +28,7 @@ require_relative 'config/model_spec'
 require_relative 'config/mcp_parser'
 
 module AIA
-  # Backward compatibility alias — existing code and tests reference AIA::ConfigSection
+  # Backward compatibility alias
   ConfigSection = MywayConfig::ConfigSection
 
   class Config < MywayConfig::Base
@@ -36,7 +36,7 @@ module AIA
     env_prefix :aia
     defaults_path File.expand_path('config/defaults.yml', __dir__)
 
-    # AIA is a CLI tool, not a Rails app — no environment sections needed.
+    # AIA is a CLI tool, not a Rails app -- no environment sections needed.
     class << self
       def validate_environment!
         # no-op: AIA has no environment sections in defaults.yml
@@ -49,7 +49,8 @@ module AIA
 
     # Nested section attributes (defined as hashes, converted to ConfigSection)
     attr_config :service, :llm, :prompts, :roles, :skills, :output, :audio, :image, :embedding,
-                :tools, :flags, :registry, :paths, :logger
+                :tools, :flags, :registry, :paths, :logger, :rules, :concurrency,
+                :tool_filter
 
     # Array/collection attributes
     attr_config :models, :pipeline, :require_libs, :mcp_servers, :mcp_use, :mcp_skip, :context_files
@@ -58,17 +59,18 @@ module AIA
     attr_accessor :prompt_id, :stdin_content, :remaining_args, :dump_file,
                   :completion, :mcp_list, :list_tools, :list_skills,
                   :executable_prompt_content,
-                  :tool_names, :loaded_tools,
+                  :tool_names, :loaded_tools, :loaded_plugins,
                   :log_level_override, :log_file_override,
-                  :connected_mcp_servers, # Array of successfully connected MCP server names
-                  :failed_mcp_servers     # Array of {name:, error:} hashes for failed MCP servers
+                  :connected_mcp_servers,  # Array of successfully connected MCP server names
+                  :mcp_server_tool_counts, # Hash of name => tool_count for connected MCP servers
+                  :failed_mcp_servers      # Array of {name:, error:} hashes for failed MCP servers
 
     # ==========================================================================
     # Type Coercion
     # ==========================================================================
 
     # Convert array of hashes to array of ModelSpec objects
-    TO_MODEL_SPECS = ->(v) {
+    TO_MODEL_SPECS = lambda { |v|
       return [] if v.nil?
       return v if v.is_a?(Array) && v.first.is_a?(ModelSpec)
 
@@ -113,8 +115,12 @@ module AIA
       embedding: config_section_coercion(:embedding),
       tools: config_section_coercion(:tools),
       flags: config_section_coercion(:flags),
+      logger: config_section_coercion(:logger),
       registry: config_section_coercion(:registry),
       paths: config_section_coercion(:paths),
+      rules: config_section_coercion(:rules),
+      concurrency: config_section_coercion(:concurrency),
+      tool_filter: config_section_coercion(:tool_filter),
 
       # Arrays
       models: TO_MODEL_SPECS,
@@ -152,50 +158,65 @@ module AIA
     # Mapping of flat CLI keys to their nested config locations
     CLI_TO_NESTED_MAP = {
       # flags section
-      chat: [:flags, :chat],
-      cost: [:flags, :cost],
-      fuzzy: [:flags, :fuzzy],
-      tokens: [:flags, :tokens],
-      no_mcp: [:flags, :no_mcp],
-      terse: [:flags, :terse],
-      debug: [:flags, :debug],
-      verbose: [:flags, :verbose],
-      consensus: [:flags, :consensus],
+      chat: %i[flags chat],
+      cost: %i[flags cost],
+      fuzzy: %i[flags fuzzy],
+      tokens: %i[flags tokens],
+      no_mcp: %i[flags no_mcp],
+      debug: %i[flags debug],
+      verbose: %i[flags verbose],
+      consensus: %i[flags consensus],
+      track_pipeline: %i[flags track_pipeline],
+      expert_routing: %i[flags expert_routing],
+      auto_tool_filter: %i[flags auto_tool_filter],
+      tool_filter_load: %i[flags tool_filter_load],
+      tool_filter_save: %i[flags tool_filter_save],
+      allow_ruby_eval: %i[flags allow_ruby_eval],
+      thinking: %i[flags thinking],
+      tool_filter_timeout_s: %i[tool_filter timeout_s],
+      concurrent_auto: %i[concurrency auto],
       # llm section
-      temperature: [:llm, :temperature],
-      max_tokens: [:llm, :max_tokens],
-      top_p: [:llm, :top_p],
-      frequency_penalty: [:llm, :frequency_penalty],
-      presence_penalty: [:llm, :presence_penalty],
+      temperature: %i[llm temperature],
+      max_tokens: %i[llm max_tokens],
+      top_p: %i[llm top_p],
+      frequency_penalty: %i[llm frequency_penalty],
+      presence_penalty: %i[llm presence_penalty],
       # prompts section
-      prompts_dir: [:prompts, :dir],
-      roles_prefix: [:prompts, :roles_prefix],
-      role: [:prompts, :role],
-      skills_dir: [:skills, :dir],
-      skills_prefix: [:prompts, :skills_prefix],
-      skills: [:prompts, :skills],
-      parameter_regex: [:prompts, :parameter_regex],
-      system_prompt: [:prompts, :system_prompt],
+      prompts_dir: %i[prompts dir],
+      roles_prefix: %i[prompts roles_prefix],
+      role: %i[prompts role],
+      skills_dir: %i[skills dir],
+      skills: %i[prompts skills],
+      tools_prefix: %i[prompts tools_prefix],
+      tool: %i[prompts tool],
+      system_prompt: %i[prompts system_prompt],
       # output section
-      output: [:output, :file],
-      history_file: [:output, :history_file],
-      append: [:output, :append],
-      markdown: [:output, :markdown],
-      # audio section
-      speak: [:audio, :speak],
-      voice: [:audio, :voice],
-      speech_model: [:audio, :speech_model],
-      transcription_model: [:audio, :transcription_model],
+      output: %i[output file],
+      history_file: %i[output history_file],
+      append: %i[output append],
+      markdown: %i[output markdown],
+      # audio section (speak is a flag, not audio config)
+      speak: %i[flags speak],
+      voice: %i[audio voice],
+      speak_command: %i[audio speak_command],
+      speech_model: %i[audio speech_model],
+      transcription_model: %i[audio transcription_model],
       # image section
-      image_size: [:image, :size],
-      image_quality: [:image, :quality],
-      image_style: [:image, :style],
+      paint: %i[flags paint],
+      image_model: %i[image model],
+      image_size: %i[image size],
+      image_quality: %i[image quality],
+      image_style: %i[image style],
       # tools section
-      tool_paths: [:tools, :paths],
-      allowed_tools: [:tools, :allowed],
-      rejected_tools: [:tools, :rejected],
+      tool_paths: %i[tools paths],
+      allowed_tools: %i[tools allowed],
+      rejected_tools: %i[tools rejected],
       # registry section
-      refresh: [:registry, :refresh],
+      refresh: %i[registry refresh],
+      # paths section
+      plugins_dir: %i[paths plugins_dir],
+      # rules section
+      rules_enabled: %i[rules enabled]
     }.freeze
 
     def initialize(overrides: {})
@@ -207,7 +228,9 @@ module AIA
       load_extra_config(extra_config_path) if extra_config_path
 
       apply_models_env_var unless overrides[:models]
+      apply_plugins_dir_env_var unless overrides[:plugins_dir]
       apply_overrides(overrides) if overrides && !overrides.empty?
+      expand_paths
       process_mcp_files(overrides[:mcp_files]) if overrides[:mcp_files]
     end
 
@@ -223,6 +246,8 @@ module AIA
           section, nested_key = CLI_TO_NESTED_MAP[key_sym]
           section_obj = send(section)
           section_obj.send("#{nested_key}=", value) if section_obj.respond_to?("#{nested_key}=")
+        elsif key_sym == :models
+          self.models = TO_MODEL_SPECS.call(Array(value))
         elsif respond_to?("#{key}=")
           send("#{key}=", value)
         elsif key.to_s.include?('__')
@@ -234,6 +259,7 @@ module AIA
     end
 
     # Convert config to hash (for dump, etc.)
+    # rubocop:disable Metrics/AbcSize
     def to_h
       {
         service: service.to_h,
@@ -248,8 +274,11 @@ module AIA
         embedding: embedding.to_h,
         tools: tools.to_h,
         flags: flags.to_h,
+        logger: logger.to_h,
         registry: registry.to_h,
         paths: paths.to_h,
+        rules: rules.to_h,
+        concurrency: concurrency.to_h,
         pipeline: pipeline,
         require_libs: require_libs,
         mcp_servers: mcp_servers,
@@ -258,41 +287,28 @@ module AIA
         context_files: context_files
       }
     end
+    # rubocop:enable Metrics/AbcSize
 
     private
 
     # Load a config file that REPLACES the user's personal config.
-    # Resets all values to bundled defaults first, then applies the
-    # file's values on top. This means -c gives you a clean slate
-    # based on defaults + the specified file only.
-    #
-    # Uses the same YAML structure as defaults.yml (nested sections).
-    # Supports both flat format and wrapped (defaults:) format.
-    #
-    # @param path [String] path to the YAML config file
     def load_extra_config(path)
       path = File.expand_path(path)
 
       unless File.exist?(path)
-        warn "ERROR: Config file not found: #{path}"
+        $stderr.puts "ERROR: Config file not found: #{path}"
         exit 1
-        return # guard for test environments that override exit
+        return # rubocop:disable Lint/UnreachableCode -- exit is mocked in tests
       end
 
-      # Reset everything to bundled defaults so the user's personal
-      # config (~/.config/aia/aia.yml) does not bleed through.
       reset_to_defaults
 
-      raw = YAML.safe_load(
-        File.read(path),
-        permitted_classes: [Symbol],
+      raw = YAML.safe_load_file(path, permitted_classes: [Symbol],
         symbolize_names: true,
-        aliases: true
-      ) || {}
+        aliases: true) || {}
 
       config_hash = raw.key?(:defaults) ? (raw[:defaults] || {}) : raw
 
-      # Store the path for reference (visible in --dump output)
       paths[:extra_config_file] = path
 
       config_hash.each do |key, value|
@@ -304,7 +320,7 @@ module AIA
         when :mcp_servers
           self.mcp_servers = Array(value)
         when :service, :llm, :prompts, :roles, :skills, :output, :audio, :image, :embedding,
-             :tools, :flags, :registry, :paths, :logger
+             :tools, :flags, :registry, :paths, :logger, :rules, :concurrency
           section = send(key)
           if section.is_a?(MywayConfig::ConfigSection) && value.is_a?(Hash)
             merge_into_section(section, value)
@@ -315,9 +331,6 @@ module AIA
       expand_paths
     end
 
-    # Reset all configuration values to the bundled defaults from
-    # defaults.yml, discarding anything set by the user's personal
-    # config file or environment variables.
     def reset_to_defaults
       defaults = self.class.schema
 
@@ -334,10 +347,6 @@ module AIA
       end
     end
 
-    # Recursively reset a ConfigSection to match the given defaults hash.
-    #
-    # @param section [MywayConfig::ConfigSection] target section
-    # @param defaults_hash [Hash] default values to restore
     def reset_section(section, defaults_hash)
       defaults_hash.each do |key, value|
         existing = section[key.to_sym]
@@ -349,11 +358,6 @@ module AIA
       end
     end
 
-    # Recursively merge a hash into a ConfigSection, preserving
-    # existing values for keys not present in the hash.
-    #
-    # @param section [MywayConfig::ConfigSection] target section
-    # @param hash [Hash] values to merge in
     def merge_into_section(section, hash)
       hash.each do |key, value|
         existing = section[key.to_sym]
@@ -365,76 +369,59 @@ module AIA
       end
     end
 
-    # Apply AIA_MODEL env var if set (supports comma-separated models with optional roles)
-    # Format: MODEL[=ROLE][,MODEL[=ROLE]]...
     def apply_models_env_var
-      models_env = ENV['AIA_MODEL']
+      models_env = ENV.fetch('AIA_MODEL', nil)
       return if models_env.nil? || models_env.empty?
 
       self.models = TO_MODEL_SPECS.call(models_env.split(',').map(&:strip))
     end
 
-    def expand_paths
-      # Expand ~ in paths
-      if paths.aia_dir
-        paths.aia_dir = File.expand_path(paths.aia_dir)
-      end
-
-      if paths.config_file
-        paths.config_file = File.expand_path(paths.config_file)
-      end
-
-      if prompts.dir
-        prompts.dir = File.expand_path(prompts.dir)
-      end
-
-      if prompts.roles_dir
-        prompts.roles_dir = File.expand_path(prompts.roles_dir)
-      end
+    # Backward-compatible top-level env var for plugin directory.
+    # Unlike most nested config values, this uses AIA_PLUGINS_DIR per requirements.
+    def apply_plugins_dir_env_var
+      plugins_dir_env = ENV.fetch('AIA_PLUGINS_DIR', nil)
+      return if plugins_dir_env.nil? || plugins_dir_env.empty?
 
-      if output.history_file
-        output.history_file = File.expand_path(output.history_file)
-      end
-
-      if roles.dir
-        roles.dir = File.expand_path(roles.dir)
-      end
+      paths.plugins_dir = plugins_dir_env
+    end
 
-      if skills.dir
-        skills.dir = File.expand_path(skills.dir)
-      end
+    def expand_paths
+      expand_path_on(paths, :aia_dir)
+      expand_path_on(paths, :config_file)
+      expand_path_on(paths, :plugins_dir)
+      expand_path_on(prompts, :dir)
+      expand_path_on(prompts, :roles_dir)
+      expand_path_on(roles, :dir)
+      expand_path_on(skills, :dir)
+      expand_path_on(tools, :dir)
+      expand_path_on(output, :history_file)
+      expand_path_on(rules, :dir)
+    end
 
-      if tools.dir
-        tools.dir = File.expand_path(tools.dir)
-      end
+    def expand_path_on(obj, method_name)
+      return unless obj.respond_to?(method_name)
+      val = obj.send(method_name)
+      return unless val
+      obj.send(:"#{method_name}=", File.expand_path(val))
     end
 
     def ensure_arrays
-      # Ensure array fields are actually arrays
       self.pipeline = [] if pipeline.nil?
       self.require_libs = [] if require_libs.nil?
       self.context_files = [] if context_files.nil?
       self.mcp_servers = [] if mcp_servers.nil?
       self.mcp_use = [] if mcp_use.nil?
       self.mcp_skip = [] if mcp_skip.nil?
-
-      # Ensure tools.paths is an array
       tools.paths = [] if tools.paths.nil?
-
-      # Ensure prompts.skills is an array
       prompts.skills = [] if prompts.respond_to?(:skills) && prompts.skills.nil?
     end
 
-    # Process MCP JSON files and merge servers into mcp_servers
-    #
-    # @param mcp_files [Array<String>] paths to MCP JSON configuration files
     def process_mcp_files(mcp_files)
       return if mcp_files.nil? || mcp_files.empty?
 
       servers_from_files = McpParser.parse_files(mcp_files)
       return if servers_from_files.empty?
 
-      # Merge with existing mcp_servers (CLI files take precedence)
       self.mcp_servers = (mcp_servers || []) + servers_from_files
     end
 
@@ -442,32 +429,10 @@ module AIA
       section = parts[0].to_sym
       key = parts[1].to_sym
 
-      case section
-      when :llm
-        llm.send("#{key}=", value) if llm.respond_to?("#{key}=")
-      when :prompts
-        prompts.send("#{key}=", value) if prompts.respond_to?("#{key}=")
-      when :output
-        output.send("#{key}=", value) if output.respond_to?("#{key}=")
-      when :audio
-        audio.send("#{key}=", value) if audio.respond_to?("#{key}=")
-      when :image
-        image.send("#{key}=", value) if image.respond_to?("#{key}=")
-      when :embedding
-        embedding.send("#{key}=", value) if embedding.respond_to?("#{key}=")
-      when :tools
-        tools.send("#{key}=", value) if tools.respond_to?("#{key}=")
-      when :flags
-        flags.send("#{key}=", value) if flags.respond_to?("#{key}=")
-      when :registry
-        registry.send("#{key}=", value) if registry.respond_to?("#{key}=")
-      when :paths
-        paths.send("#{key}=", value) if paths.respond_to?("#{key}=")
-      when :roles
-        roles.send("#{key}=", value) if roles.respond_to?("#{key}=")
-      when :skills
-        skills.send("#{key}=", value) if skills.respond_to?("#{key}=")
-      end
+      target = respond_to?(section) ? send(section) : nil
+      return unless target.respond_to?("#{key}=")
+
+      target.send("#{key}=", value)
     end
   end
 end

data/lib/aia/content_extractor.rb

@@ -0,0 +1,153 @@
+# frozen_string_literal: true
+
+# lib/aia/content_extractor.rb
+#
+# Shared module for extracting text content from various response types.
+# Handles RobotLab::RobotResult, SimpleFlow::Result, and plain strings.
+# Include this module instead of defining extract_content locally.
+
+module AIA
+  module ContentExtractor
+    # Extract text content from a response object.
+    #
+    # @param result the response (RobotResult, SimpleFlow::Result, String, etc.)
+    # @return [String] the extracted text
+    def extract_content(result)
+      if defined?(SimpleFlow::Result) && result.is_a?(SimpleFlow::Result)
+        return extract_network_content(result)
+      end
+
+      if result.respond_to?(:reply)
+        result.reply
+      elsif result.respond_to?(:last_text_content)
+        result.last_text_content
+      elsif result.respond_to?(:content)
+        result.content
+      else
+        result.to_s
+      end
+    end
+
+    # Extract content from a Network's SimpleFlow::Result.
+    # Each robot's result is stored in context under its task name.
+    # Includes per-robot timing when available.
+    #
+    # @param flow_result [SimpleFlow::Result]
+    # @return [String]
+    def extract_network_content(flow_result)
+      parts = []
+      flow_result.context.each do |task_name, robot_result|
+        next if task_name == :run_params
+
+        content = if robot_result.respond_to?(:reply)
+                    robot_result.reply
+                  elsif robot_result.respond_to?(:content)
+                    robot_result.content
+                  else
+                    robot_result.to_s
+                  end
+        next unless content && !content.empty?
+
+        duration = robot_result.respond_to?(:duration) ? robot_result.duration : nil
+        robot_name = robot_result.respond_to?(:robot_name) ? robot_result.robot_name : nil
+        label = robot_name && robot_name != task_name.to_s ? "#{robot_name} [#{task_name}]" : task_name.to_s
+        header = if duration
+                   "**#{label}** (#{format_duration(duration)}):"
+                 else
+                   "**#{label}:**"
+                 end
+        parts << "#{header}\n#{content}"
+      end
+      parts.join("\n\n")
+    end
+
+    # Store structured robot results into a network's shared memory.
+    # Each robot's output is written as a hash under result_<task_name>.
+    #
+    # @param flow_result [SimpleFlow::Result]
+    # @param network [RobotLab::Network]
+    def store_results_in_memory(flow_result, network)
+      return unless network.respond_to?(:memory)
+
+      memory = network.memory
+      flow_result.context.each do |task_name, robot_result|
+        next if task_name == :run_params
+        next unless robot_result.respond_to?(:reply)
+
+        memory.current_writer = task_name.to_s
+        memory.set(:"result_#{task_name}", {
+          content:  robot_result.reply,
+          model:    robot_result.respond_to?(:robot_name) ? robot_result.robot_name : nil,
+          duration: robot_result.respond_to?(:duration) ? robot_result.duration : nil
+        })
+      end
+    end
+
+    # Present an AI result: extract content, record the turn, display,
+    # write to output file, show metrics, speak, and print separator.
+    # Callers pass only the keyword args they need; everything optional
+    # defaults to nil / false so this works for both full ChatLoop turns
+    # and simpler SpecialModeHandler dispatches.
+    #
+    # @param result the response object
+    # @param streamed_content [String, nil] pre-streamed content (skips display_ai_response)
+    # @param prompt [String, nil] user prompt for tracker
+    # @param elapsed [Float, nil] elapsed seconds
+    # @param ui_presenter [UIPresenter] for display calls
+    # @param tracker [SessionTracker, nil] to record the turn
+    # @param decisions [Hash, nil] routing decisions for tracker
+    # @return [String] the extracted content
+    def present_result(result, ui_presenter:, streamed_content: nil, prompt: nil, elapsed: nil, tracker: nil, decisions: nil)
+      content = streamed_content || extract_content(result)
+
+      if tracker && prompt
+        tracker.record_turn(
+          model: AIA.config.models.first.name,
+          input: prompt,
+          result: result,
+          decisions: decisions,
+          elapsed: elapsed
+        )
+      end
+
+      if streamed_content
+        puts
+      else
+        ui_presenter.display_ai_response(content)
+      end
+
+      output_to_file(content)
+      display_metrics(result, elapsed: elapsed) if respond_to?(:display_metrics, true)
+      speak(content) if respond_to?(:speak, true)
+      ui_presenter.display_separator
+
+      content
+    end
+
+    # Write AI response content to the configured output file.
+    #
+    # @param content [String] the AI response text
+    def output_to_file(content)
+      out_file = AIA.config.output.file
+      return unless out_file
+
+      File.open(out_file, 'a') { |f| f.puts "\nAI: #{content}" }
+    end
+
+    # Format a duration in seconds to a human-readable string.
+    #
+    # @param seconds [Float, nil]
+    # @return [String]
+    def format_duration(seconds)
+      return "0.0s" unless seconds
+
+      if seconds < 60
+        "%.1fs" % seconds
+      else
+        minutes = (seconds / 60).to_i
+        secs = seconds % 60
+        "#{minutes}m %04.1fs" % secs
+      end
+    end
+  end
+end

data/lib/aia/cost_calculator.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+# lib/aia/cost_calculator.rb
+#
+# Shared cost calculation logic used by SessionTracker and UIPresenter.
+# Looks up per-million token pricing from RubyLLM::Models and returns
+# a breakdown hash.
+
+module AIA
+  module CostCalculator
+    # @param model_id [String, nil]
+    # @param input_tokens [Integer]
+    # @param output_tokens [Integer]
+    # @return [Hash] { available: true, total_cost:, input_cost:, output_cost: }
+    #             or { available: false } when pricing is unavailable
+    def self.calculate(model_id:, input_tokens:, output_tokens:)
+      return { available: false } unless model_id && defined?(RubyLLM::Models)
+
+      # RubyLLM::Models.find raises ModelNotFoundError on an unknown id (it never
+      # returns nil), so the unknown-model case is handled by the rescue below.
+      model_info = RubyLLM::Models.find(model_id)
+
+      input_price  = model_info.respond_to?(:input_price_per_million)  ? model_info.input_price_per_million  : nil
+      output_price = model_info.respond_to?(:output_price_per_million) ? model_info.output_price_per_million : nil
+      return { available: false } unless input_price && output_price
+
+      input_cost  = input_tokens  * input_price  / 1_000_000.0
+      output_cost = output_tokens * output_price / 1_000_000.0
+
+      { available: true, total_cost: input_cost + output_cost, input_cost: input_cost, output_cost: output_cost }
+    rescue RubyLLM::ModelNotFoundError
+      # Expected when pricing for the model isn't in the registry — not an error.
+      { available: false }
+    rescue StandardError => e
+      { available: false, error: e.message }
+    end
+  end
+end