ai-agents 0.4.3 → 0.6.0

data/lib/agents/runner.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require_relative "message_extractor"
-
 module Agents
   # The execution engine that orchestrates conversations between users and agents.
   # Runner manages the conversation flow, handles tool execution through RubyLLM,
@@ -55,6 +53,7 @@ module Agents
     DEFAULT_MAX_TURNS = 10
 
     class MaxTurnsExceeded < StandardError; end
+    class AgentNotFoundError < StandardError; end
 
     # Create a thread-safe agent runner for multi-agent conversations.
     # The first agent becomes the default entry point for new conversations.
@@ -79,9 +78,10 @@ module Agents
     # @param context [Hash] Shared context data accessible to all tools
     # @param registry [Hash] Registry of agents for handoff resolution
     # @param max_turns [Integer] Maximum conversation turns before stopping
+    # @param headers [Hash, nil] Custom HTTP headers passed to the underlying LLM provider
     # @param callbacks [Hash] Optional callbacks for real-time event notifications
     # @return [RunResult] The result containing output, messages, and usage
-    def run(starting_agent, input, context: {}, registry: {}, max_turns: DEFAULT_MAX_TURNS, callbacks: {})
+    def run(starting_agent, input, context: {}, registry: {}, max_turns: DEFAULT_MAX_TURNS, headers: nil, callbacks: {})
       # The starting_agent is already determined by AgentRunner based on conversation history
       current_agent = starting_agent
 
@@ -90,15 +90,22 @@ module Agents
       context_wrapper = RunContext.new(context_copy, callbacks: callbacks)
       current_turn = 0
 
+      runtime_headers = Helpers::Headers.normalize(headers)
+      agent_headers = Helpers::Headers.normalize(current_agent.headers)
+
       # Create chat and restore conversation history
-      chat = create_chat(current_agent, context_wrapper)
+      chat = RubyLLM::Chat.new(model: current_agent.model)
+      current_headers = Helpers::Headers.merge(agent_headers, runtime_headers)
+      apply_headers(chat, current_headers)
+      configure_chat_for_agent(chat, current_agent, context_wrapper, replace: false)
       restore_conversation_history(chat, context_wrapper)
 
+
       loop do
         current_turn += 1
         raise MaxTurnsExceeded, "Exceeded maximum turns: #{max_turns}" if current_turn > max_turns
 
-        # Get response from LLM (Extended Chat handles tool execution with handoff detection)
+        # Get response from LLM (RubyLLM handles tool execution with halting based handoff detection)
         result = if current_turn == 1
                    # Emit agent thinking event for initial message
                    context_wrapper.callback_manager.emit_agent_thinking(current_agent.name, input)
@@ -110,17 +117,23 @@ module Agents
                  end
         response = result
 
-        # Check for handoff response from our extended chat
-        if response.is_a?(Agents::Chat::HandoffResponse)
-          next_agent = response.target_agent
+        # Check for handoff via RubyLLM's halt mechanism
+        if response.is_a?(RubyLLM::Tool::Halt) && context_wrapper.context[:pending_handoff]
+          handoff_info = context_wrapper.context.delete(:pending_handoff)
+          next_agent = handoff_info[:target_agent]
 
           # Validate that the target agent is in our registry
           # This prevents handoffs to agents that weren't explicitly provided
           unless registry[next_agent.name]
-            puts "[Agents] Warning: Handoff to unregistered agent '#{next_agent.name}', continuing with current agent"
-            next if response.tool_call?
-
-            next
+            save_conversation_state(chat, context_wrapper, current_agent)
+            error = AgentNotFoundError.new("Handoff failed: Agent '#{next_agent.name}' not found in registry")
+            return RunResult.new(
+              output: nil,
+              messages: Helpers::MessageExtractor.extract_messages(chat, current_agent),
+              usage: context_wrapper.usage,
+              context: context_wrapper.context,
+              error: error
+            )
           end
 
           # Save current conversation state before switching
@@ -133,9 +146,11 @@ module Agents
           current_agent = next_agent
           context_wrapper.context[:current_agent] = next_agent.name
 
-          # Create new chat for new agent with restored history
-          chat = create_chat(current_agent, context_wrapper)
-          restore_conversation_history(chat, context_wrapper)
+          # Reconfigure existing chat for new agent - preserves conversation history automatically
+          configure_chat_for_agent(chat, current_agent, context_wrapper, replace: true)
+          agent_headers = Helpers::Headers.normalize(current_agent.headers)
+          current_headers = Helpers::Headers.merge(agent_headers, runtime_headers)
+          apply_headers(chat, current_headers)
 
           # Force the new agent to respond to the conversation context
           # This ensures the user gets a response from the new agent
@@ -143,6 +158,17 @@ module Agents
           next
         end
 
+        # Handle non-handoff halts - return the halt content as final response
+        if response.is_a?(RubyLLM::Tool::Halt)
+          save_conversation_state(chat, context_wrapper, current_agent)
+          return RunResult.new(
+            output: response.content,
+            messages: Helpers::MessageExtractor.extract_messages(chat, current_agent),
+            usage: context_wrapper.usage,
+            context: context_wrapper.context
+          )
+        end
+
         # If tools were called, continue the loop to let them execute
         next if response.tool_call?
 
@@ -153,7 +179,7 @@ module Agents
 
         return RunResult.new(
           output: response.content,
-          messages: MessageExtractor.extract_messages(chat, current_agent),
+          messages: Helpers::MessageExtractor.extract_messages(chat, current_agent),
           usage: context_wrapper.usage,
           context: context_wrapper.context
         )
@@ -164,7 +190,7 @@ module Agents
 
       RunResult.new(
         output: "Conversation ended: #{e.message}",
-        messages: chat ? MessageExtractor.extract_messages(chat, current_agent) : [],
+        messages: chat ? Helpers::MessageExtractor.extract_messages(chat, current_agent) : [],
         usage: context_wrapper.usage,
         error: e,
         context: context_wrapper.context
@@ -175,7 +201,7 @@ module Agents
 
       RunResult.new(
         output: nil,
-        messages: chat ? MessageExtractor.extract_messages(chat, current_agent) : [],
+        messages: chat ? Helpers::MessageExtractor.extract_messages(chat, current_agent) : [],
         usage: context_wrapper.usage,
         error: e,
         context: context_wrapper.context
@@ -184,6 +210,11 @@ module Agents
 
     private
 
+    # Creates a deep copy of context data for thread safety.
+    # Preserves conversation history array structure while avoiding agent mutation.
+    #
+    # @param context [Hash] The context to copy
+    # @return [Hash] Thread-safe deep copy of the context
     def deep_copy_context(context)
       # Handle deep copying for thread safety
       context.dup.tap do |copied|
@@ -194,31 +225,40 @@ module Agents
       end
     end
 
+    # Restores conversation history from context into RubyLLM chat.
+    # Converts stored message hashes back into RubyLLM::Message objects with proper content handling.
+    #
+    # @param chat [RubyLLM::Chat] The chat instance to restore history into
+    # @param context_wrapper [RunContext] Context containing conversation history
     def restore_conversation_history(chat, context_wrapper)
       history = context_wrapper.context[:conversation_history] || []
 
       history.each do |msg|
         # Only restore user and assistant messages with content
         next unless %i[user assistant].include?(msg[:role].to_sym)
-        next unless msg[:content] && !MessageExtractor.content_empty?(msg[:content])
+        next unless msg[:content] && !Helpers::MessageExtractor.content_empty?(msg[:content])
 
-        chat.add_message(
+        # Extract text content safely - handle both string and hash content
+        content = RubyLLM::Content.new(msg[:content])
+
+        # Create a proper RubyLLM::Message and pass it to add_message
+        message = RubyLLM::Message.new(
           role: msg[:role].to_sym,
-          content: msg[:content]
+          content: content
         )
-      rescue StandardError => e
-        # Continue with partial history on error
-        puts "[Agents] Failed to restore message: #{e.message}"
+        chat.add_message(message)
       end
-    rescue StandardError => e
-      # If history restoration completely fails, continue with empty history
-      puts "[Agents] Failed to restore conversation history: #{e.message}"
-      context_wrapper.context[:conversation_history] = []
     end
 
+    # Saves current conversation state from RubyLLM chat back to context for persistence.
+    # Maintains conversation continuity across agent handoffs and process boundaries.
+    #
+    # @param chat [RubyLLM::Chat] The chat instance to extract state from
+    # @param context_wrapper [RunContext] Context to save state into
+    # @param current_agent [Agents::Agent] The currently active agent
     def save_conversation_state(chat, context_wrapper, current_agent)
       # Extract messages from chat
-      messages = MessageExtractor.extract_messages(chat, current_agent)
+      messages = Helpers::MessageExtractor.extract_messages(chat, current_agent)
 
       # Update context with latest state
       context_wrapper.context[:conversation_history] = messages
@@ -230,29 +270,59 @@ module Agents
       context_wrapper.context.delete(:pending_handoff)
     end
 
-    def create_chat(agent, context_wrapper)
+    # Configures a RubyLLM chat instance with agent-specific settings.
+    # Uses RubyLLM's replace option to swap agent context while preserving conversation history during handoffs.
+    #
+    # @param chat [RubyLLM::Chat] The chat instance to configure
+    # @param agent [Agents::Agent] The agent whose configuration to apply
+    # @param context_wrapper [RunContext] Thread-safe context wrapper
+    # @param replace [Boolean] Whether to replace existing configuration (true for handoffs, false for initial setup)
+    # @return [RubyLLM::Chat] The configured chat instance
+    def configure_chat_for_agent(chat, agent, context_wrapper, replace: false)
       # Get system prompt (may be dynamic)
       system_prompt = agent.get_system_prompt(context_wrapper)
 
-      # Separate handoff tools from regular tools
-      handoff_tools = agent.handoff_agents.map { |target_agent| HandoffTool.new(target_agent) }
-      regular_tools = agent.tools
+      # Combine all tools - both handoff and regular tools need wrapping
+      all_tools = build_agent_tools(agent, context_wrapper)
 
-      # Only wrap regular tools - handoff tools will be handled directly by Chat
-      wrapped_regular_tools = regular_tools.map { |tool| ToolWrapper.new(tool, context_wrapper) }
+      # Switch model if different (important for handoffs between agents using different models)
+      chat.with_model(agent.model) if replace
 
-      # Create extended chat with handoff awareness and context
-      chat = Agents::Chat.new(
-        model: agent.model,
-        temperature: agent.temperature,
-        handoff_tools: handoff_tools,        # Direct tools, no wrapper
-        context_wrapper: context_wrapper,    # Pass context directly
-        response_schema: agent.response_schema # Pass structured output schema
-      )
+      # Configure chat with instructions, temperature, tools, and schema
+      chat.with_instructions(system_prompt, replace: replace) if system_prompt
+      chat.with_temperature(agent.temperature) if agent.temperature
+      chat.with_tools(*all_tools, replace: replace)
+      chat.with_schema(agent.response_schema) if agent.response_schema
 
-      chat.with_instructions(system_prompt) if system_prompt
-      chat.with_tools(*wrapped_regular_tools) if wrapped_regular_tools.any?
       chat
     end
+
+    def apply_headers(chat, headers)
+      return if headers.empty?
+
+      chat.with_headers(**headers)
+    end
+
+    # Builds thread-safe tool wrappers for an agent's tools and handoff tools.
+    #
+    # @param agent [Agents::Agent] The agent whose tools to wrap
+    # @param context_wrapper [RunContext] Thread-safe context wrapper for tool execution
+    # @return [Array<ToolWrapper>] Array of wrapped tools ready for RubyLLM
+    def build_agent_tools(agent, context_wrapper)
+      all_tools = []
+
+      # Add handoff tools
+      agent.handoff_agents.each do |target_agent|
+        handoff_tool = HandoffTool.new(target_agent)
+        all_tools << ToolWrapper.new(handoff_tool, context_wrapper)
+      end
+
+      # Add regular tools
+      agent.tools.each do |tool|
+        all_tools << ToolWrapper.new(tool, context_wrapper)
+      end
+
+      all_tools
+    end
   end
 end

data/lib/agents/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Agents
-  VERSION = "0.4.3"
+  VERSION = "0.6.0"
 end

data/lib/agents.rb

@@ -111,12 +111,11 @@ require_relative "agents/run_context"
 require_relative "agents/tool_context"
 require_relative "agents/tool"
 require_relative "agents/handoff"
+require_relative "agents/helpers"
 require_relative "agents/agent"
 
 # Execution components
-require_relative "agents/chat"
 require_relative "agents/tool_wrapper"
-require_relative "agents/message_extractor"
 require_relative "agents/callback_manager"
 require_relative "agents/agent_runner"
 require_relative "agents/runner"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ai-agents
 version: !ruby/object:Gem::Version
-  version: 0.4.3
+  version: 0.6.0
 platform: ruby
 authors:
 - Shivam Mishra
@@ -15,14 +15,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '1.3'
+        version: 1.8.2
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '1.3'
+        version: 1.8.2
 description: Ruby AI Agents SDK enables creating complex AI workflows with multi-agent
   orchestration, tool execution, safety guardrails, and provider-agnostic LLM integration.
 email:
@@ -60,6 +60,7 @@ files:
 - docs/guides/agent-as-tool-pattern.md
 - docs/guides/multi-agent-systems.md
 - docs/guides/rails-integration.md
+- docs/guides/request-headers.md
 - docs/guides/state-persistence.md
 - docs/guides/structured-output.md
 - docs/index.md
@@ -101,9 +102,10 @@ files:
 - lib/agents/agent_runner.rb
 - lib/agents/agent_tool.rb
 - lib/agents/callback_manager.rb
-- lib/agents/chat.rb
 - lib/agents/handoff.rb
-- lib/agents/message_extractor.rb
+- lib/agents/helpers.rb
+- lib/agents/helpers/headers.rb
+- lib/agents/helpers/message_extractor.rb
 - lib/agents/result.rb
 - lib/agents/run_context.rb
 - lib/agents/runner.rb

data/lib/agents/chat.rb

@@ -1,161 +0,0 @@
-# frozen_string_literal: true
-
-require_relative "tool_context"
-
-module Agents
-  # Extended chat class that inherits from RubyLLM::Chat but adds proper handoff handling.
-  # This solves the infinite handoff loop problem by treating handoffs as turn-ending
-  # operations rather than allowing auto-continuation.
-  class Chat < RubyLLM::Chat
-    # Response object that indicates a handoff occurred
-    class HandoffResponse
-      attr_reader :target_agent, :response, :handoff_message
-
-      def initialize(target_agent:, response:, handoff_message:)
-        @target_agent = target_agent
-        @response = response
-        @handoff_message = handoff_message
-      end
-
-      def tool_call?
-        true
-      end
-
-      def content
-        @handoff_message
-      end
-    end
-
-    def initialize(model: nil, handoff_tools: [], context_wrapper: nil, temperature: nil, response_schema: nil,
-                   **options)
-      super(model: model, **options)
-      @handoff_tools = handoff_tools
-      @context_wrapper = context_wrapper
-
-      # Set temperature if provided (RubyLLM::Chat sets this via accessor)
-      @temperature = temperature if temperature
-
-      # Set response schema if provided
-      with_schema(response_schema) if response_schema
-
-      # Register handoff tools with RubyLLM for schema generation
-      @handoff_tools.each { |tool| with_tool(tool) }
-    end
-
-    # Override the problematic auto-execution method from RubyLLM::Chat
-    def complete(&block)
-      @on[:new_message]&.call
-      response = @provider.complete(
-        messages,
-        tools: @tools,
-        temperature: @temperature,
-        model: @model.id,
-        connection: @connection,
-        params: @params,
-        schema: @schema,
-        &block
-      )
-      @on[:end_message]&.call(response)
-
-      # Handle JSON parsing for structured output (like RubyLLM::Chat)
-      if @schema && response.content.is_a?(String)
-        begin
-          response.content = JSON.parse(response.content)
-        rescue JSON::ParserError
-          # If parsing fails, keep content as string
-        end
-      end
-
-      add_message(response)
-
-      if response.tool_call?
-        handle_tools_with_handoff_detection(response, &block)
-      else
-        response
-      end
-    end
-
-    private
-
-    def handle_tools_with_handoff_detection(response, &block)
-      handoff_calls, regular_calls = classify_tool_calls(response.tool_calls)
-
-      if handoff_calls.any?
-        # Execute first handoff only
-        handoff_result = execute_handoff_tool(handoff_calls.first)
-
-        # Add tool result to conversation
-        add_tool_result(handoff_calls.first.id, handoff_result[:message])
-
-        # Return handoff response to signal agent switch (ends turn)
-        HandoffResponse.new(
-          target_agent: handoff_result[:target_agent],
-          response: response,
-          handoff_message: handoff_result[:message]
-        )
-      else
-        # Use RubyLLM's original tool execution for regular tools
-        execute_regular_tools_and_continue(regular_calls, &block)
-      end
-    end
-
-    def classify_tool_calls(tool_calls)
-      handoff_tool_names = @handoff_tools.map(&:name).map(&:to_s)
-
-      handoff_calls = []
-      regular_calls = []
-
-      tool_calls.each_value do |tool_call|
-        if handoff_tool_names.include?(tool_call.name)
-          handoff_calls << tool_call
-        else
-          regular_calls << tool_call
-        end
-      end
-
-      [handoff_calls, regular_calls]
-    end
-
-    def execute_handoff_tool(tool_call)
-      tool = @handoff_tools.find { |t| t.name.to_s == tool_call.name }
-      raise "Handoff tool not found: #{tool_call.name}" unless tool
-
-      # Execute the handoff tool directly with context
-      tool_context = ToolContext.new(run_context: @context_wrapper)
-      result = tool.execute(tool_context, **{}) # Handoff tools take no additional params
-
-      {
-        target_agent: tool.target_agent,
-        message: result.to_s
-      }
-    end
-
-    def execute_regular_tools_and_continue(tool_calls, &block)
-      # Execute each regular tool call
-      tool_calls.each do |tool_call|
-        @on[:new_message]&.call
-        result = execute_tool(tool_call)
-        message = add_tool_result(tool_call.id, result)
-        @on[:end_message]&.call(message)
-      end
-
-      # Continue conversation after tool execution
-      complete(&block)
-    end
-
-    # Reuse RubyLLM's existing tool execution logic
-    def execute_tool(tool_call)
-      tool = tools[tool_call.name.to_sym]
-      args = tool_call.arguments
-      tool.call(args)
-    end
-
-    def add_tool_result(tool_use_id, result)
-      add_message(
-        role: :tool,
-        content: result.is_a?(Hash) && result[:error] ? result[:error] : result.to_s,
-        tool_call_id: tool_use_id
-      )
-    end
-  end
-end

data/lib/agents/message_extractor.rb

@@ -1,97 +0,0 @@
-# frozen_string_literal: true
-
-module Agents
-  # Service object responsible for extracting and formatting conversation messages
-  # from RubyLLM chat objects into a format suitable for persistence and context restoration.
-  #
-  # Handles different message types:
-  # - User messages: Basic content preservation
-  # - Assistant messages: Includes agent attribution and tool calls
-  # - Tool result messages: Links back to original tool calls
-  #
-  # @example Extract messages from a chat
-  #   messages = MessageExtractor.extract_messages(chat, current_agent)
-  #   #=> [
-  #     { role: :user, content: "Hello" },
-  #     { role: :assistant, content: "Hi!", agent_name: "Support", tool_calls: [...] },
-  #     { role: :tool, content: "Result", tool_call_id: "call_123" }
-  #   ]
-  class MessageExtractor
-    # Check if content is considered empty (handles both String and Hash content)
-    #
-    # @param content [String, Hash, nil] The content to check
-    # @return [Boolean] true if content is empty, false otherwise
-    def self.content_empty?(content)
-      case content
-      when String
-        content.strip.empty?
-      when Hash
-        content.empty?
-      else
-        content.nil?
-      end
-    end
-
-    # Extract messages from a chat object for conversation history persistence
-    #
-    # @param chat [Object] Chat object that responds to :messages
-    # @param current_agent [Agent] The agent currently handling the conversation
-    # @return [Array<Hash>] Array of message hashes suitable for persistence
-    def self.extract_messages(chat, current_agent)
-      new(chat, current_agent).extract
-    end
-
-    def initialize(chat, current_agent)
-      @chat = chat
-      @current_agent = current_agent
-    end
-
-    def extract
-      return [] unless @chat.respond_to?(:messages)
-
-      @chat.messages.filter_map do |msg|
-        case msg.role
-        when :user, :assistant
-          extract_user_or_assistant_message(msg)
-        when :tool
-          extract_tool_message(msg)
-        end
-      end
-    end
-
-    private
-
-    def extract_user_or_assistant_message(msg)
-      return nil unless msg.content && !self.class.content_empty?(msg.content)
-
-      message = {
-        role: msg.role,
-        content: msg.content
-      }
-
-      if msg.role == :assistant
-        # Add agent attribution for conversation continuity
-        message[:agent_name] = @current_agent.name if @current_agent
-
-        # Add tool calls if present
-        if msg.tool_call? && msg.tool_calls
-          # RubyLLM stores tool_calls as Hash with call_id => ToolCall object
-          # Reference: RubyLLM::StreamAccumulator#tool_calls_from_stream
-          message[:tool_calls] = msg.tool_calls.values.map(&:to_h)
-        end
-      end
-
-      message
-    end
-
-    def extract_tool_message(msg)
-      return nil unless msg.tool_result?
-
-      {
-        role: msg.role,
-        content: msg.content,
-        tool_call_id: msg.tool_call_id
-      }
-    end
-  end
-end