aia 0.9.8 → 0.9.9

data/lib/aia/context_manager.rb

@@ -26,14 +26,9 @@ module AIA
     # @param system_prompt [String, nil] The system prompt to potentially prepend.
     # @return [Array<Hash>] The conversation context array.
     def get_context(system_prompt: nil)
-      # Ensure system prompt is present if provided and not already the first message
-      if  system_prompt &&
-          !system_prompt.strip.empty? &&
-          (
-            @context.empty? ||
-            @context.first[:role] != 'system'
-          )
-         add_system_prompt(system_prompt)
+      # Add or replace system prompt if provided and not empty
+      if system_prompt && !system_prompt.strip.empty?
+        add_system_prompt(system_prompt)
       end
       @context
     end

data/lib/aia/directive_processor.rb

@@ -162,19 +162,19 @@ module AIA
       spaces  = " "*indent
       width   = TTY::Screen.width - indent - 2
 
-      if !AIA.config.tools.empty?
+      if AIA.config.tools.empty?
+        puts "No tools are available"
+      else
         puts
         puts "Available Tools"
         puts "==============="
 
-        AIA.config.tools.split(',').map(&:strip).each do |tool|
-          klass = tool.constantize
-          puts "\n#{klass.name}"
-          puts "-"*klass.name.size
-          puts WordWrapper::MinimumRaggedness.new(width, klass.description).wrap.split("\n").map{|s| spaces+s+"\n"}.join
+        AIA.config.tools.each do |tool|
+          name = tool.respond_to?(:name) ? tool.name : tool.class.name
+          puts "\n#{name}"
+          puts "-"*name.size
+          puts WordWrapper::MinimumRaggedness.new(width, tool.description).wrap.split("\n").map{|s| spaces+s+"\n"}.join
         end
-      else
-        puts "No tools configured"
       end
       puts
 
@@ -260,7 +260,15 @@ module AIA
 
     desc "Shortcut for //config model _and_ //config model = value"
     def model(args, context_manager=nil)
-      send(:config, args.prepend('model'), context_manager)
+      if args.empty?
+        puts
+        puts AIA.config.client.model.to_h.pretty_inspect
+        puts
+      else
+        send(:config, args.prepend('model'), context_manager)
+      end
+
+      return ''
     end
 
     desc "Shortcut for //config temperature _and_ //config temperature = value"
@@ -345,9 +353,12 @@ module AIA
       counter = 0
 
       RubyLLM.models.all.each do |llm|
+        cw      = llm.context_window
+        caps    = llm.capabilities.join(',')
         inputs  = llm.modalities.input.join(',')
         outputs = llm.modalities.output.join(',')
-        entry   = "- #{llm.id} (#{llm.provider}) #{inputs} to #{outputs}"
+        mode    = "#{inputs} to #{outputs}"
+        entry   = "- #{llm.id} (#{llm.provider}) cw: #{cw} mode: #{mode} caps: #{caps}"
 
         if query.nil? || query.empty?
           counter += 1

data/lib/aia/ruby_llm_adapter.rb

@@ -66,31 +66,81 @@ module AIA
       end
     end
 
+
     def setup_chat_with_tools
       begin
-        @chat = RubyLLM.chat(model: @model)
+        @chat  = RubyLLM.chat(model: @model)
+        @model = @chat.model.name if @model.nil? # using default model
       rescue => e
         STDERR.puts "ERROR: #{e.message}"
         exit 1
       end
 
-      return unless @chat.model.supports_functions?
+      unless @chat.model.supports_functions?
+        AIA.config.tools      = []
+        AIA.config.tool_names = ""
+        return
+      end
+
+      load_tools
 
-      if !AIA.config.tool_paths.empty? && !@chat.model.supports?(:function_calling)
-        STDERR.puts "ERROR: The model #{@model} does not support tools"
-        exit 1
+      @chat.with_tools(*tools) unless tools.empty?
+    end
+
+
+    def load_tools
+      @tools = []
+
+      support_local_tools
+      support_mcp
+      filter_tools_by_allowed_list
+      filter_tools_by_rejected_list
+      drop_duplicate_tools
+
+      if tools.empty?
+        AIA.config.tool_names = ""
+      else
+        AIA.config.tool_names = @tools.map(&:name).join(', ')
+        AIA.config.tools      = @tools
       end
+    end
+
 
-      @tools = ObjectSpace.each_object(Class).select do |klass|
+    def support_local_tools
+      @tools += ObjectSpace.each_object(Class).select do |klass|
         klass < RubyLLM::Tool
       end
+    end
+
 
-      unless tools.empty?
-        @chat.with_tools(*tools)
-        AIA.config.tools = tools.map(&:name).join(', ')
+    def support_mcp
+      RubyLLM::MCP.establish_connection
+      @tools += RubyLLM::MCP.tools
+    rescue => e
+      STDERR.puts "Warning: Failed to connect MCP clients: #{e.message}"
+    end
+
+
+    def drop_duplicate_tools
+      seen_names = Set.new
+      original_size = @tools.size
+
+      @tools.select! do |tool|
+        tool_name = tool.name
+        if seen_names.include?(tool_name)
+          STDERR.puts "WARNING: Duplicate tool name detected: '#{tool_name}'. Only the first occurrence will be used."
+          false
+        else
+          seen_names.add(tool_name)
+          true
+        end
       end
+
+      removed_count = original_size - @tools.size
+      STDERR.puts "Removed #{removed_count} duplicate tools" if removed_count > 0
     end
 
+
     # TODO: Need to rethink this dispatcher pattern w/r/t RubyLLM's capabilities
     #       This code was originally designed for AiClient
     #
@@ -117,7 +167,7 @@ module AIA
     end
 
     def transcribe(audio_file)
-      @chat.ask("Transcribe this audio", with: audio_file)
+      @chat.ask("Transcribe this audio", with: audio_file).content
     end
 
     def speak(text)
@@ -195,6 +245,24 @@ module AIA
 
     private
 
+    def filter_tools_by_allowed_list
+      return if AIA.config.allowed_tools.nil?
+
+      @tools.select! do |tool|
+        tool_name = tool.respond_to?(:name) ? tool.name : tool.class.name
+        AIA.config.allowed_tools.any? { |allowed| tool_name.include?(allowed) }
+      end
+    end
+
+    def filter_tools_by_rejected_list
+      return if AIA.config.rejected_tools.nil?
+
+      @tools.reject! do |tool|
+        tool_name = tool.respond_to?(:name) ? tool.name : tool.class.name
+        AIA.config.rejected_tools.any? { |rejected| tool_name.include?(rejected) }
+      end
+    end
+
     def extract_model_parts
       parts = AIA.config.model.split('/')
       parts.map!(&:strip)

data/lib/aia/session.rb

@@ -22,26 +22,36 @@ module AIA
 
     def initialize(prompt_handler)
       @prompt_handler = prompt_handler
-      @chat_prompt_id = nil  # Initialize to nil
-      @include_context_flag = true  # Flag to include context files, stdin, etc. only once
+      @chat_prompt_id = nil
+      @include_context_flag = true
+      
+      setup_prompt_and_history_manager
+      initialize_components
+      setup_output_file
+    end
 
+    def setup_prompt_and_history_manager
       # Special handling for chat mode with context files but no prompt ID
-      if AIA.chat? && AIA.config.prompt_id.empty? && AIA.config.context_files && !AIA.config.context_files.empty?
+      if AIA.chat? && (AIA.config.prompt_id.nil? || AIA.config.prompt_id.empty?) && AIA.config.context_files && !AIA.config.context_files.empty?
         prompt_instance = nil
         @history_manager = nil
-      elsif AIA.chat? && AIA.config.prompt_id.empty?
+      elsif AIA.chat? && (AIA.config.prompt_id.nil? || AIA.config.prompt_id.empty?)
         prompt_instance = nil
         @history_manager = nil
       else
         prompt_instance = @prompt_handler.get_prompt(AIA.config.prompt_id)
         @history_manager = HistoryManager.new(prompt: prompt_instance)
       end
+    end
 
+    def initialize_components
       @context_manager = ContextManager.new(system_prompt: AIA.config.system_prompt)
       @ui_presenter = UIPresenter.new
       @directive_processor = DirectiveProcessor.new
       @chat_processor = ChatProcessorService.new(@ui_presenter, @directive_processor)
+    end
 
+    def setup_output_file
       if AIA.config.out_file && !AIA.config.out_file.nil? && !AIA.append? && File.exist?(AIA.config.out_file)
         File.open(AIA.config.out_file, "w") { } # Truncate the file
       end
@@ -95,37 +105,44 @@ module AIA
 
     # Process a single prompt with all its requirements
     def process_single_prompt(prompt_id)
-      role_id = AIA.config.role
-
       # Skip empty prompt IDs
       return if prompt_id.nil? || prompt_id.empty?
 
-      # Get and process the prompt
+      prompt = setup_prompt_processing(prompt_id)
+      return unless prompt
+
+      prompt_text = finalize_prompt_text(prompt)
+      send_prompt_and_get_response(prompt_text)
+    end
+
+    def setup_prompt_processing(prompt_id)
+      role_id = AIA.config.role
+
       begin
         prompt = @prompt_handler.get_prompt(prompt_id, role_id)
       rescue StandardError => e
         puts "Error processing prompt '#{prompt_id}': #{e.message}"
-        return
+        return nil
       end
 
-      # Collect variable values (only when flag is true)
       if @include_context_flag
         collect_variable_values(prompt)
         enhance_prompt_with_extras(prompt)
       end
 
-      # Get final prompt text
+      prompt
+    end
+
+    def finalize_prompt_text(prompt)
       prompt_text = prompt.to_s
 
-      # Add context files (only when flag is true)
       if @include_context_flag
         prompt_text = add_context_files(prompt_text)
         # SMELL: TODO? empty the AIA.config.context_files array
         @include_context_flag = false
       end
 
-      # Send prompt to AI and get response
-      send_prompt_and_get_response(prompt_text)
+      prompt_text
     end
 
     # Collect variable values from user input
@@ -144,17 +161,19 @@ module AIA
           history_values: history,
         )
 
-        # Update history
-        history.delete(value) if history.include?(value)
-        history << value
-        history.shift if history.size > HistoryManager::MAX_VARIABLE_HISTORY
-
-        variable_values[var_name] = history
+        variable_values[var_name] = update_variable_history(history, value)
       end
 
       prompt.parameters = variable_values
     end
 
+    def update_variable_history(history, value)
+      history.delete(value) if history.include?(value)
+      history << value
+      history.shift if history.size > HistoryManager::MAX_VARIABLE_HISTORY
+      history
+    end
+
     # Add terse instructions, stdin content, and executable prompt file content
     def enhance_prompt_with_extras(prompt)
       # Add terse instruction if needed
@@ -210,157 +229,187 @@ module AIA
     # NOTE: there could have been an initial prompt sent into this session
     #       via a prompt_id on the command line, piped in text, or context files.
     def start_chat(skip_context_files: false)
+      setup_chat_session
+      process_initial_context(skip_context_files)
+      handle_piped_input
+      run_chat_loop
+    ensure
+      @ui_presenter.display_chat_end
+    end
+
+    private
+
+    def setup_chat_session
+      initialize_chat_ui
+      @chat_prompt_id = generate_chat_prompt_id
+      create_temporary_prompt
+      setup_signal_handlers
+      create_chat_prompt_object
+      Reline::HISTORY.clear
+    end
+
+    def initialize_chat_ui
       puts "\nEntering interactive chat mode..."
       @ui_presenter.display_chat_header
+    end
 
-      # Generate chat prompt ID
+    def generate_chat_prompt_id
       now = Time.now
-      @chat_prompt_id = "chat_#{now.strftime("%Y%m%d_%H%M%S")}"
+      "chat_#{now.strftime("%Y%m%d_%H%M%S")}"
+    end
 
-      # Create the temporary prompt
-      begin
-        # Create the unique? prompt ID in the file storage system with its initial text
-        PromptManager::Prompt.create(
-          id: @chat_prompt_id,
-          text: "Today's date is #{now.strftime("%Y-%m-%d")} and the current time is #{now.strftime("%H:%M:%S")}",
-        )
+    def create_temporary_prompt
+      now = Time.now
+      PromptManager::Prompt.create(
+        id: @chat_prompt_id,
+        text: "Today's date is #{now.strftime("%Y-%m-%d")} and the current time is #{now.strftime("%H:%M:%S")}",
+      )
+    end
 
-        # Capture self for the handlers
-        session_instance = self
-
-        # Set up cleanup handlers only after prompt is created
-        at_exit { session_instance.send(:cleanup_chat_prompt) }
-        Signal.trap("INT") {
-          session_instance.send(:cleanup_chat_prompt)
-          exit
-        }
-
-        # Access this chat session's prompt object in order to do the dynamic things
-        # in follow up prompts that can be done in the batch mode like shell substitution. etc.
-        @chat_prompt = PromptManager::Prompt.new(
-          id: @chat_prompt_id,
-          directives_processor: @directive_processor,
-          erb_flag: true,
-          envar_flag: true,
-          external_binding: binding,
-        )
+    def setup_signal_handlers
+      session_instance = self
+      at_exit { session_instance.send(:cleanup_chat_prompt) }
+      Signal.trap("INT") {
+        session_instance.send(:cleanup_chat_prompt)
+        exit
+      }
+    end
 
-        Reline::HISTORY.clear
+    def create_chat_prompt_object
+      @chat_prompt = PromptManager::Prompt.new(
+        id: @chat_prompt_id,
+        directives_processor: @directive_processor,
+        erb_flag: true,
+        envar_flag: true,
+        external_binding: binding,
+      )
+    end
 
-        # Load context files if any and not skipping
-        if !skip_context_files && AIA.config.context_files && !AIA.config.context_files.empty?
-          context = AIA.config.context_files.map do |file|
-            File.read(file) rescue "Error reading file: #{file}"
-          end.join("\n\n")
+    def process_initial_context(skip_context_files)
+      return if skip_context_files || !AIA.config.context_files || AIA.config.context_files.empty?
 
-          if !context.empty?
-            # Add context files content to context
-            @context_manager.add_to_context(role: "user", content: context)
+      context = AIA.config.context_files.map do |file|
+        File.read(file) rescue "Error reading file: #{file}"
+      end.join("\n\n")
 
-            # Process the context
-            @ui_presenter.display_thinking_animation
-            response = @chat_processor.process_prompt(@context_manager.get_context)
+      return if context.empty?
 
-            # Add AI response to context
-            @context_manager.add_to_context(role: "assistant", content: response)
+      # Add context files content to context
+      @context_manager.add_to_context(role: "user", content: context)
 
-            # Output the response
-            @chat_processor.output_response(response)
-            @chat_processor.speak(response)
-            @ui_presenter.display_separator
-          end
-        end
+      # Process the context
+      @ui_presenter.display_thinking_animation
+      response = @chat_processor.process_prompt(@context_manager.get_context)
+
+      # Add AI response to context
+      @context_manager.add_to_context(role: "assistant", content: response)
+
+      # Output the response
+      @chat_processor.output_response(response)
+      @chat_processor.speak(response)
+      @ui_presenter.display_separator
+    end
 
-        # Handle piped input
-        if !STDIN.tty?
-          original_stdin = STDIN.dup
-          piped_input = STDIN.read.strip
-          STDIN.reopen("/dev/tty")
+    def handle_piped_input
+      return if STDIN.tty?
 
-          if !piped_input.empty?
-            @chat_prompt.text = piped_input
-            processed_input = @chat_prompt.to_s
+      original_stdin = STDIN.dup
+      piped_input = STDIN.read.strip
+      STDIN.reopen("/dev/tty")
 
-            @context_manager.add_to_context(role: "user", content: processed_input)
+      return if piped_input.empty?
 
-            @ui_presenter.display_thinking_animation
-            response = @chat_processor.process_prompt(@context_manager.get_context)
+      @chat_prompt.text = piped_input
+      processed_input = @chat_prompt.to_s
+
+      @context_manager.add_to_context(role: "user", content: processed_input)
+
+      @ui_presenter.display_thinking_animation
+      response = @chat_processor.process_prompt(@context_manager.get_context)
+
+      @context_manager.add_to_context(role: "assistant", content: response)
+      @chat_processor.output_response(response)
+      @chat_processor.speak(response) if AIA.speak?
+      @ui_presenter.display_separator
+
+      STDIN.reopen(original_stdin)
+    end
 
-            @context_manager.add_to_context(role: "assistant", content: response)
-            @chat_processor.output_response(response)
-            @chat_processor.speak(response) if AIA.speak?
-            @ui_presenter.display_separator
+    def run_chat_loop
+      loop do
+        follow_up_prompt = @ui_presenter.ask_question
+
+        break if follow_up_prompt.nil? || follow_up_prompt.strip.downcase == "exit" || follow_up_prompt.strip.empty?
+
+        if AIA.config.out_file
+          File.open(AIA.config.out_file, "a") do |file|
+            file.puts "\nYou: #{follow_up_prompt}"
           end
+        end
 
-          STDIN.reopen(original_stdin)
+        if @directive_processor.directive?(follow_up_prompt)
+          follow_up_prompt = process_chat_directive(follow_up_prompt)
+          next if follow_up_prompt.nil?
         end
 
-        # Main chat loop
-        loop do
-          follow_up_prompt = @ui_presenter.ask_question
+        @chat_prompt.text = follow_up_prompt
+        processed_prompt = @chat_prompt.to_s
 
-          break if follow_up_prompt.nil? || follow_up_prompt.strip.downcase == "exit" || follow_up_prompt.strip.empty?
+        @context_manager.add_to_context(role: "user", content: processed_prompt)
+        conversation = @context_manager.get_context
 
-          if AIA.config.out_file
-            File.open(AIA.config.out_file, "a") do |file|
-              file.puts "\nYou: #{follow_up_prompt}"
-            end
-          end
+        @ui_presenter.display_thinking_animation
+        response = @chat_processor.process_prompt(conversation)
 
-          if @directive_processor.directive?(follow_up_prompt)
-            directive_output = @directive_processor.process(follow_up_prompt, @context_manager)
-
-            if follow_up_prompt.strip.start_with?("//clear")
-              # The directive processor has called context_manager.clear_context
-              # but we need a more aggressive approach to fully clear all context
-
-              # First, clear the context manager's context
-              @context_manager.clear_context(keep_system_prompt: true)
-
-              # Second, try clearing the client's context
-              if AIA.config.client && AIA.config.client.respond_to?(:clear_context)
-                AIA.config.client.clear_context
-              end
-
-              # Third, completely reinitialize the client to ensure fresh state
-              # This is the most aggressive approach to ensure no context remains
-              begin
-                AIA.config.client = AIA::RubyLLMAdapter.new
-              rescue => e
-                STDERR.puts "Error reinitializing client: #{e.message}"
-              end
-
-              @ui_presenter.display_info("Chat context cleared.")
-              next
-            elsif directive_output.nil? || directive_output.strip.empty?
-              next
-            else
-              puts "\n#{directive_output}\n"
-              follow_up_prompt = "I executed this directive: #{follow_up_prompt}\nHere's the output: #{directive_output}\nLet's continue our conversation."
-            end
-          end
+        @ui_presenter.display_ai_response(response)
+        @context_manager.add_to_context(role: "assistant", content: response)
+        @chat_processor.speak(response)
+
+        @ui_presenter.display_separator
+      end
+    end
 
-          @chat_prompt.text = follow_up_prompt
-          processed_prompt = @chat_prompt.to_s
+    def process_chat_directive(follow_up_prompt)
+      directive_output = @directive_processor.process(follow_up_prompt, @context_manager)
+      
+      return handle_clear_directive if follow_up_prompt.strip.start_with?("//clear")
+      return handle_empty_directive_output if directive_output.nil? || directive_output.strip.empty?
+      
+      handle_successful_directive(follow_up_prompt, directive_output)
+    end
 
-          @context_manager.add_to_context(role: "user", content: processed_prompt)
-          conversation = @context_manager.get_context
+    def handle_clear_directive
+      # The directive processor has called context_manager.clear_context
+      # but we need a more aggressive approach to fully clear all context
 
-          @ui_presenter.display_thinking_animation
-          response = @chat_processor.process_prompt(conversation)
+      # First, clear the context manager's context
+      @context_manager.clear_context(keep_system_prompt: true)
 
-          @ui_presenter.display_ai_response(response)
-          @context_manager.add_to_context(role: "assistant", content: response)
-          @chat_processor.speak(response)
+      # Second, try clearing the client's context
+      if AIA.config.client && AIA.config.client.respond_to?(:clear_context)
+        AIA.config.client.clear_context
+      end
 
-          @ui_presenter.display_separator
-        end
-      ensure
-        @ui_presenter.display_chat_end
+      # Third, completely reinitialize the client to ensure fresh state
+      # This is the most aggressive approach to ensure no context remains
+      begin
+        AIA.config.client = AIA::RubyLLMAdapter.new
+      rescue => e
+        STDERR.puts "Error reinitializing client: #{e.message}"
       end
+
+      @ui_presenter.display_info("Chat context cleared.")
+      nil
     end
 
-    private
+    def handle_empty_directive_output
+      nil
+    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."
+    end
 
     def cleanup_chat_prompt
       if @chat_prompt_id

data/lib/aia/ui_presenter.rb

@@ -54,9 +54,9 @@ module AIA
         line = line.chomp
 
         # Check for code block delimiters
-        if line.match?(/^```(\w*)$/) && !in_code_block
+        if (match = line.match(/^```(\w*)$/)) && !in_code_block
           in_code_block = true
-          language = $1
+          language = match[1]
           output.puts "#{indent}```#{language}"
         elsif line.match?(/^```$/) && in_code_block
           in_code_block = false
@@ -107,9 +107,11 @@ module AIA
         spinner = TTY::Spinner.new("[:spinner] #{spinner_message}", format: :bouncing_ball)
         spinner.auto_spin
 
-        result = yield
-
-        spinner.stop
+        begin
+          result = yield
+        ensure
+          spinner.stop
+        end
         result
       else
         yield