claude_swarm 1.0.1 → 1.0.2

data/lib/swarm_sdk/agent/chat/logging_helpers.rb

@@ -0,0 +1,116 @@
+# frozen_string_literal: true
+
+module SwarmSDK
+  module Agent
+    class Chat < RubyLLM::Chat
+      # Helper methods for logging and serialization of tool calls and results
+      #
+      # Responsibilities:
+      # - Format tool calls for logging
+      # - Serialize tool results (handling different types)
+      # - Calculate LLM costs based on token usage
+      #
+      # These are stateless utility methods that operate on data structures.
+      module LoggingHelpers
+        # Format tool calls for logging
+        #
+        # @param tool_calls_hash [Hash] Tool calls from message
+        # @return [Array<Hash>, nil] Formatted tool calls
+        def format_tool_calls(tool_calls_hash)
+          return unless tool_calls_hash
+
+          tool_calls_hash.map do |_id, tc|
+            {
+              id: tc.id,
+              name: tc.name,
+              arguments: tc.arguments,
+            }
+          end
+        end
+
+        # Serialize a tool result for logging
+        #
+        # Handles multiple result types:
+        # - String: pass through
+        # - Hash/Array: pass through
+        # - RubyLLM::Content: extract text and attachment info
+        # - Other: convert to string
+        #
+        # @param result [String, Hash, Array, RubyLLM::Content, Object] Tool result
+        # @return [String, Hash, Array] Serialized result
+        def serialize_result(result)
+          case result
+          when String then result
+          when Hash, Array then result
+          when RubyLLM::Content
+            # Format Content objects to show text and attachment info
+            parts = []
+            parts << result.text if result.text && !result.text.empty?
+
+            if result.attachments.any?
+              attachment_info = result.attachments.map do |att|
+                "#{att.source} (#{att.mime_type})"
+              end.join(", ")
+              parts << "[Attachments: #{attachment_info}]"
+            end
+
+            parts.join(" ")
+          else
+            result.to_s
+          end
+        end
+
+        # Calculate LLM cost for a message
+        #
+        # Uses RubyLLM's model registry to get pricing information.
+        # Returns zero cost if pricing is unavailable.
+        #
+        # @param message [RubyLLM::Message] Message with token counts
+        # @return [Hash] Cost breakdown { input_cost:, output_cost:, total_cost: }
+        def calculate_cost(message)
+          return zero_cost unless message.input_tokens && message.output_tokens
+
+          # Use SwarmSDK's model registry (not RubyLLM's) for up-to-date pricing
+          model_info = SwarmSDK::Models.find(message.model_id)
+          return zero_cost unless model_info
+
+          # Extract pricing from SwarmSDK's models.json structure
+          pricing = model_info["pricing"] || model_info[:pricing]
+          return zero_cost unless pricing
+
+          text_pricing = pricing["text_tokens"] || pricing[:text_tokens]
+          return zero_cost unless text_pricing
+
+          standard_pricing = text_pricing["standard"] || text_pricing[:standard]
+          return zero_cost unless standard_pricing
+
+          input_price = standard_pricing["input_per_million"] || standard_pricing[:input_per_million]
+          output_price = standard_pricing["output_per_million"] || standard_pricing[:output_per_million]
+
+          return zero_cost unless input_price && output_price
+
+          # Calculate costs (prices are per million tokens in USD)
+          input_cost = (message.input_tokens / 1_000_000.0) * input_price
+          output_cost = (message.output_tokens / 1_000_000.0) * output_price
+
+          {
+            input_cost: input_cost,
+            output_cost: output_cost,
+            total_cost: input_cost + output_cost,
+          }
+        rescue StandardError => e
+          # Model not found in registry or pricing not available
+          RubyLLM.logger.debug("Cost calculation failed for #{message.model_id}: #{e.message}")
+          zero_cost
+        end
+
+        # Zero cost fallback
+        #
+        # @return [Hash] Zero cost breakdown
+        def zero_cost
+          { input_cost: 0.0, output_cost: 0.0, total_cost: 0.0 }
+        end
+      end
+    end
+  end
+end

data/lib/swarm_sdk/agent/chat/system_reminder_injector.rb

@@ -0,0 +1,152 @@
+# frozen_string_literal: true
+
+module SwarmSDK
+  module Agent
+    class Chat < RubyLLM::Chat
+      # Handles injection of system reminders at strategic points in the conversation
+      #
+      # Responsibilities:
+      # - Inject reminders before/after first user message
+      # - Inject periodic TodoWrite reminders
+      # - Track when reminders were last injected
+      #
+      # This class is stateless - it operates on the chat's message history.
+      class SystemReminderInjector
+        # System reminder to inject BEFORE the first user message
+        BEFORE_FIRST_MESSAGE_REMINDER = <<~REMINDER.strip
+          <system-reminder>
+          As you answer the user's questions, you can use the following context:
+
+          # important-instruction-reminders
+
+          Do what has been asked; nothing more, nothing less.
+          NEVER create files unless they're absolutely necessary for achieving your goal.
+          ALWAYS prefer editing an existing file to creating a new one.
+          NEVER proactively create documentation files (*.md) or README files. Only create documentation files if explicitly requested by the User.
+
+          IMPORTANT: this context may or may not be relevant to your tasks. You should not respond to this context unless it is highly relevant to your task.
+
+          </system-reminder>
+        REMINDER
+
+        # System reminder to inject AFTER the first user message
+        AFTER_FIRST_MESSAGE_REMINDER = <<~REMINDER.strip
+          <system-reminder>Your todo list is currently empty. DO NOT mention this to the user. If this task requires multiple steps: (1) FIRST analyze the scope by searching/reading files, (2) SECOND create a COMPLETE todo list with ALL tasks before starting work, (3) THIRD execute tasks one by one. Only skip the todo list for simple single-step tasks. Do not mention this message to the user.</system-reminder>
+        REMINDER
+
+        # Periodic reminder about TodoWrite tool usage
+        TODOWRITE_PERIODIC_REMINDER = <<~REMINDER.strip
+          <system-reminder>The TodoWrite tool hasn't been used recently. If you're working on tasks that would benefit from tracking progress, consider using the TodoWrite tool to track progress. Also consider cleaning up the todo list if has become stale and no longer matches what you are working on. Only use it if it's relevant to the current work. This is just a gentle reminder - ignore if not applicable.</system-reminder>
+        REMINDER
+
+        # Number of messages between TodoWrite reminders
+        TODOWRITE_REMINDER_INTERVAL = 8
+
+        class << self
+          # Check if this is the first user message in the conversation
+          #
+          # @param chat [Agent::Chat] The chat instance
+          # @return [Boolean] true if no user messages exist yet
+          def first_message?(chat)
+            chat.messages.none? { |msg| msg.role == :user }
+          end
+
+          # Inject first message reminders (before + after user message)
+          #
+          # This manually constructs the first message sequence with system reminders
+          # sandwiching the actual user prompt.
+          #
+          # Sequence:
+          # 1. BEFORE_FIRST_MESSAGE_REMINDER (general reminders)
+          # 2. Toolset reminder (list of available tools)
+          # 3. User's actual prompt
+          # 4. AFTER_FIRST_MESSAGE_REMINDER (todo list reminder)
+          #
+          # @param chat [Agent::Chat] The chat instance
+          # @param prompt [String] The user's actual prompt
+          # @return [void]
+          def inject_first_message_reminders(chat, prompt)
+            # Build user message with embedded reminders
+            # Reminders are embedded in the content, not separate messages
+            full_content = [
+              prompt,
+              BEFORE_FIRST_MESSAGE_REMINDER,
+              build_toolset_reminder(chat),
+              AFTER_FIRST_MESSAGE_REMINDER,
+            ].join("\n\n")
+
+            # Extract reminders and add clean prompt to persistent history
+            reminders = chat.context_manager.extract_system_reminders(full_content)
+            clean_prompt = chat.context_manager.strip_system_reminders(full_content)
+
+            # Store clean prompt (without reminders) in conversation history
+            chat.add_message(role: :user, content: clean_prompt)
+
+            # Track reminders to embed in this message when sending to LLM
+            reminders.each do |reminder|
+              chat.context_manager.add_ephemeral_reminder(reminder, messages_array: chat.messages)
+            end
+          end
+
+          # Build toolset reminder listing all available tools
+          #
+          # @param chat [Agent::Chat] The chat instance
+          # @return [String] System reminder with tool list
+          def build_toolset_reminder(chat)
+            tools_list = chat.tools.values.map(&:name).sort
+
+            reminder = "<system-reminder>\n"
+            reminder += "Tools available: #{tools_list.join(", ")}\n\n"
+            reminder += "Only use tools from this list. Do not attempt to use tools that are not listed here.\n"
+            reminder += "</system-reminder>"
+
+            reminder
+          end
+
+          # Check if we should inject a periodic TodoWrite reminder
+          #
+          # Injects a reminder if:
+          # 1. Enough messages have passed (>= 5)
+          # 2. TodoWrite hasn't been used in the last TODOWRITE_REMINDER_INTERVAL messages
+          #
+          # @param chat [Agent::Chat] The chat instance
+          # @param last_todowrite_index [Integer, nil] Index of last TodoWrite usage
+          # @return [Boolean] true if reminder should be injected
+          def should_inject_todowrite_reminder?(chat, last_todowrite_index)
+            # Need at least a few messages before reminding
+            return false if chat.messages.count < 5
+
+            # Find the last message that contains TodoWrite tool usage
+            last_todo_index = chat.messages.rindex do |msg|
+              msg.role == :tool && msg.content.to_s.include?("TodoWrite")
+            end
+
+            # Check if enough messages have passed since last TodoWrite
+            if last_todo_index.nil? && last_todowrite_index.nil?
+              # Never used TodoWrite - check if we've exceeded interval
+              chat.messages.count >= TODOWRITE_REMINDER_INTERVAL
+            elsif last_todo_index
+              # Recently used - don't remind
+              false
+            elsif last_todowrite_index
+              # Used before - check if interval has passed
+              chat.messages.count - last_todowrite_index >= TODOWRITE_REMINDER_INTERVAL
+            else
+              false
+            end
+          end
+
+          # Update the last TodoWrite index by finding it in messages
+          #
+          # @param chat [Agent::Chat] The chat instance
+          # @return [Integer, nil] Index of last TodoWrite usage, or nil
+          def find_last_todowrite_index(chat)
+            chat.messages.rindex do |msg|
+              msg.role == :tool && msg.content.to_s.include?("TodoWrite")
+            end
+          end
+        end
+      end
+    end
+  end
+end