llm_gateway 0.6.0 → 0.7.0

data/lib/llm_gateway/adapters/structs.rb

@@ -52,17 +52,43 @@ class ToolCall < BaseStruct
   end
 end
 
+class ServerToolCall < ToolCall
+  attribute :type, Types::String.enum("server_tool_use")
+end
+
 class ToolResult < BaseStruct
-  attribute :type, Types::String.enum("tool_result")
+  attribute :type, Types::String
   attribute :tool_use_id, Types::String
-  attribute :content, Types::String
+  attribute :content, Types::Any
+
+  def to_h
+    {
+      type: type,
+      tool_use_id: tool_use_id,
+      content: content
+    }
+  end
+end
+
+class ServerToolResult < ToolResult
+  attribute :type, Types::String.enum("server_tool_result")
+  attribute? :name, Types::String.optional
+
+  def to_h
+    result = super
+    result[:name] = name unless name.nil?
+    result
+  end
 end
 
 class PartialAssistantMessage < BaseStruct
   ContentBlock =
     Types.Instance(TextContent) |
     Types.Instance(ReasoningContent) |
-    Types.Instance(ToolCall)
+    Types.Instance(ToolCall) |
+    Types.Instance(ServerToolCall) |
+    Types.Instance(ToolResult) |
+    Types.Instance(ServerToolResult)
 
   attribute? :id, Types::String.optional
   attribute? :model, Types::String.optional
@@ -78,7 +104,7 @@ class PartialAssistantMessage < BaseStruct
   end
 
   def self.build_content_block(block)
-    return block if block.is_a?(TextContent) || block.is_a?(ReasoningContent) || block.is_a?(ToolCall)
+    return block if block.is_a?(TextContent) || block.is_a?(ReasoningContent) || block.is_a?(ToolCall) || block.is_a?(ServerToolCall) || block.is_a?(ToolResult) || block.is_a?(ServerToolResult)
 
     case block[:type] || block["type"]
     when "text"
@@ -86,14 +112,16 @@ class PartialAssistantMessage < BaseStruct
     when "reasoning"
       ReasoningContent.new(block)
     when "thinking"
-      ReasoningContent.new(
-        type: "reasoning",
-        reasoning: block[:thinking] || block["thinking"] || block[:reasoning] || block["reasoning"],
-        signature: block[:signature] || block["signature"]
-      )
+      ReasoningContent.new(type: "reasoning", reasoning: block[:thinking] || block["thinking"] || block[:reasoning] || block["reasoning"], signature: block[:signature] || block["signature"])
     when "tool_use"
       ToolCall.new(block)
+    when "server_tool_use"
+      ServerToolCall.new(block)
     else
+      type = block[:type] || block["type"]
+      return ServerToolResult.new(block) if type == "server_tool_result"
+      return ToolResult.new(block) if type&.end_with?("_tool_result")
+
       raise ArgumentError, "Unsupported content block type: #{block[:type] || block['type']}"
     end
   end
@@ -102,7 +130,7 @@ class PartialAssistantMessage < BaseStruct
 end
 
 class AssistantStreamEvent < BaseStruct
-  EventType = Types::Coercible::Symbol.enum(:text_start, :text_delta, :text_end, :tool_start, :tool_delta, :tool_end, :reasoning_start, :reasoning_delta, :reasoning_end)
+  EventType = Types::Coercible::Symbol.enum(:text_start, :text_delta, :text_end, :tool_start, :tool_delta, :tool_end, :tool_result_start, :tool_result_delta, :tool_result_end, :reasoning_start, :reasoning_delta, :reasoning_end)
 
   attribute :type, EventType
   attribute :delta, Types::Coercible::String.default { "" }
@@ -144,6 +172,13 @@ end
 class AssistantToolStartEvent < AssistantStreamEvent
   attribute :id, Types::String
   attribute :name, Types::String
+  attribute :tool_type, Types::String.default("tool_use".freeze).enum("tool_use", "server_tool_use")
+  attribute :content_index, Types::Integer
+end
+
+class AssistantToolResultStartEvent < AssistantStreamEvent
+  attribute :tool_use_id, Types::String
+  attribute :name, Types::String
   attribute :content_index, Types::Integer
 end
 

data/lib/llm_gateway/agents/event.rb

@@ -0,0 +1,105 @@
+# frozen_string_literal: true
+
+require_relative "../adapters/structs"
+
+module LlmGateway
+  module Agents
+    module Event
+      AgentEventType = Types::Coercible::Symbol.enum(
+        :agent_start,
+        :turn_start,
+        :message_start,
+        :message_update,
+        :message_end,
+        :tool_execution_start,
+        :tool_execution_end,
+        :turn_end,
+        :agent_end
+      )
+
+      StreamEvent =
+        Types.Instance(AssistantStreamEvent) |
+        Types.Instance(AssistantStreamMessageEvent) |
+        Types.Instance(AssistantStreamMessageEndEvent)
+
+      ToolParameters = Types::Hash.schema(
+        id: Types::String,
+        type: Types::String.enum("tool_use"),
+        name: Types::String,
+        input: Types::Hash
+      )
+
+      class ToolCallResult < ::BaseStruct
+        attribute :type, Types::Coercible::Symbol.enum(:tool_result)
+        attribute :tool_use_id, Types::String
+        attribute :content, Types::Any
+
+        def to_h
+          {
+            type: type.to_s,
+            tool_use_id: tool_use_id,
+            content: content
+          }
+        end
+
+        def dig(*keys)
+          to_h.dig(*keys)
+        end
+      end
+
+      class Base < ::BaseStruct
+        attribute :type, AgentEventType
+
+        def to_h
+          {
+            type: type
+          }
+        end
+      end
+
+      class AgentStart < Base
+        attribute :type, Types::Coercible::Symbol.default(:agent_start).enum(:agent_start)
+      end
+
+      class TurnStart < Base
+        attribute :type, Types::Coercible::Symbol.default(:turn_start).enum(:turn_start)
+      end
+
+      class MessageStart < Base
+        attribute :type, Types::Coercible::Symbol.default(:message_start).enum(:message_start)
+      end
+
+      class MessageUpdate < Base
+        attribute :type, Types::Coercible::Symbol.default(:message_update).enum(:message_update)
+        attribute :stream_event, StreamEvent
+      end
+
+      class MessageEnd < Base
+        attribute :type, Types::Coercible::Symbol.default(:message_end).enum(:message_end)
+        attribute :message, Types.Instance(AssistantMessage)
+      end
+
+      class ToolExecutionStart < Base
+        attribute :type, Types::Coercible::Symbol.default(:tool_execution_start).enum(:tool_execution_start)
+        attribute :parameters, ToolParameters
+      end
+
+      class ToolExecutionEnd < Base
+        attribute :type, Types::Coercible::Symbol.default(:tool_execution_end).enum(:tool_execution_end)
+        attribute :parameters, ToolParameters
+        attribute :result, ToolCallResult
+      end
+
+      class TurnEnd < Base
+        attribute :type, Types::Coercible::Symbol.default(:turn_end).enum(:turn_end)
+        attribute :message, Types.Instance(AssistantMessage)
+        attribute :tool_results, Types::Array.of(Types.Instance(::ToolResult))
+      end
+
+      class AgentEnd < Base
+        attribute :type, Types::Coercible::Symbol.default(:agent_end).enum(:agent_end)
+        attribute :messages, Types::Array.of(Types::Hash)
+      end
+    end
+  end
+end

data/lib/llm_gateway/agents/file_session_manager.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+require "fileutils"
+require "json"
+require "securerandom"
+require "time"
+require_relative "in_memory_session_manager"
+
+module LlmGateway
+  module Agents
+    class FileSessionManager < InMemorySessionManager
+      attr_reader :file_name, :session_path
+
+      def initialize(file_name = nil, session_id: nil, session_start: nil, session_dir: nil)
+        super(session_id)
+        @file_name = file_name
+        @preset_session_start = session_start
+        @session_dir = session_dir
+      end
+
+      def session_id
+        events
+        @session_id
+      end
+
+      def session_start
+        events
+        @session_start
+      end
+
+      def normalize_path(file_name)
+        File.expand_path(file_name)
+      end
+
+      def events
+        @events ||= begin
+          @session_path = normalize_path(file_name) if file_name
+          if @session_path && File.exist?(@session_path)
+            load_session(@session_path)
+          else
+            create_new_session
+          end
+        end
+      end
+
+      private
+
+      def create_new_session
+        @session_id ||= SecureRandom.uuid
+        @session_start = @preset_session_start || Time.now.strftime("%Y%m%d_%H%M%S")
+
+        session_event = {
+          type: "session",
+          id: @session_id,
+          timestamp: @session_start
+        }
+
+        @session_path ||= File.join(session_dir, "#{@session_start}_#{@session_id}.jsonl")
+        FileUtils.mkdir_p(File.dirname(@session_path))
+        File.open(@session_path, "a") do |file|
+          file.puts(JSON.generate(session_event))
+        end
+
+        [ session_event ]
+      end
+
+      def load_session(path)
+        loaded_events = []
+        File.foreach(path).with_index(1) do |line, line_number|
+          next if line.strip.empty?
+
+          loaded_events << JSON.parse(line, symbolize_names: true)
+        rescue JSON::ParserError => e
+          raise ArgumentError, "Invalid JSONL in #{path} at line #{line_number}: #{e.message}"
+        end
+
+        session_event = loaded_events.find { |event| event[:type] == "session" }
+        @session_id = session_event[:id] if session_event&.dig(:id)
+        @session_start = session_event[:timestamp] if session_event&.dig(:timestamp)
+
+        loaded_events
+      end
+
+      def persist_entry(entry)
+        attributes = super
+
+        FileUtils.mkdir_p(File.dirname(session_path))
+        File.open(session_path, "a") do |file|
+          file.puts(JSON.generate(entry))
+        end
+
+        attributes
+      end
+
+      def session_dir
+        File.expand_path(@session_dir || ENV.fetch("LLM_GATEWAY_SESSION_DIR", "~/.llm_gateway/sessions"))
+      end
+    end
+  end
+end

data/lib/llm_gateway/agents/harness.rb

@@ -0,0 +1,176 @@
+# frozen_string_literal: true
+
+require_relative "event"
+require_relative "../utils"
+
+module LlmGateway
+  module Agents
+    class Harness < LlmGateway::Prompt
+      COMPACTION_TOKEN_THRESHOLD = 180_000
+      COMPACTION_IDLE_THRESHOLD_SECONDS = 60 * 60
+      attr_accessor :provider
+      attr_reader :session_manager, :default_queue_mode, :queue_drain_mode,
+                  :model, :reasoning
+
+      def initialize(session_manager, provider:, model: nil, reasoning: "high")
+        @provider = provider
+        super(provider: provider, model: model, reasoning: reasoning)
+        @session_manager = session_manager
+        sync_initial_configuration_events
+        self.default_queue_mode = :next_turn
+        self.queue_drain_mode = :all
+      end
+
+      def transcript
+        session_manager.build_model_input_messages
+      end
+      alias :prompt :transcript
+
+      def prompt_message(message, &block)
+        enqueue_or_run_message(message, default_queue_mode, &block)
+      end
+
+      def steer_message(message, &block)
+        enqueue_or_run_message(message, :steer, &block)
+      end
+
+      def follow_up_message(message, &block)
+        enqueue_or_run_message(message, :follow_up, &block)
+      end
+
+      def next_turn_message(message, &block)
+        enqueue_or_run_message(message, :next_turn, &block)
+      end
+
+      def default_queue_mode=(mode)
+        @default_queue_mode = session_manager.validate_queue!(mode)
+      end
+
+      def queue_drain_mode=(mode)
+        @queue_drain_mode = session_manager.validate_drain_mode!(mode)
+      end
+
+      def model=(model_id)
+        return @model if @model == model_id
+
+        @model = model_id
+        publish_session_event(type: "model_change", model_id: model_id)
+        @model
+      end
+
+      def reasoning=(level)
+        return @reasoning if @reasoning == level
+
+        @reasoning = level
+        publish_session_event(type: "reasoning_change", reasoning: level)
+        @reasoning
+      end
+
+      def compact
+        session_manager.compaction(provider)
+      end
+
+      def run(&block)
+        emit(Event::AgentStart.new, &block)
+        drain_queue(:steer)
+        emit(Event::TurnStart.new, &block)
+        emit(Event::MessageStart.new, &block)
+
+        assistant_message = stream do |event|
+          emit(Event::MessageUpdate.new(stream_event: event), &block)
+        end
+
+        session_manager.push_message(assistant_message.to_h)
+        emit(Event::MessageEnd.new(message: assistant_message), &block)
+
+        tool_results = tool_requests(assistant_message).map do |message|
+          parameters = message.to_h
+          emit(Event::ToolExecutionStart.new(parameters: parameters), &block)
+          tool_result = find_and_execute_tool(message)
+          emit(Event::ToolExecutionEnd.new(parameters: parameters, result: tool_result), &block)
+          tool_result
+        end
+
+        tool_result_content = tool_results.map(&:to_h)
+        session_manager.push_message(
+          role: "user",
+          content: tool_result_content,
+        ) unless tool_result_content.empty?
+
+        turn_end_event = Event::TurnEnd.new(message: assistant_message, tool_results: tool_results)
+        emit(turn_end_event, &block)
+
+        if tool_results.length.positive?
+          return run(&block)
+        end
+
+        if session_manager.queued_messages?(:follow_up)
+          compact_if_needed
+          return run(&block) if drain_queue(:follow_up).any?
+        end
+
+        emit(Event::AgentEnd.new(messages: []), &block)
+        assistant_message
+      end
+      alias :continue   :run
+
+      private
+
+      def publish_session_event(type:, **attributes)
+        session_manager.push_entry(type: type, **attributes)
+      end
+
+      def sync_initial_configuration_events
+        publish_session_event(type: "model_change", model_id: model) if model && !session_manager.last_model_used
+        if reasoning && !session_manager.last_reasoning_level_used
+          publish_session_event(type: "reasoning_change", reasoning: reasoning)
+        end
+      end
+
+      def enqueue_or_run_message(message, queue, &block)
+        prepared_input = LlmGateway::Utils.deep_symbolize_keys(message)
+        result = session_manager.start_or_enqueue_user_message(prepared_input, queue: queue) do
+          compact_if_needed
+        end
+        return if result == session_manager.class::MESSAGE_QUEUED
+
+        begin
+          continue(&block)
+
+          loop do
+            break unless session_manager.queued_messages?(:next_turn)
+
+            compact_if_needed
+            drain_queue(:next_turn)
+            continue(&block)
+          end
+        ensure
+          session_manager.idle!
+        end
+      end
+
+      def compact_if_needed
+        compact if compaction_needed?
+      end
+
+      def compaction_needed?
+        session_manager.total_tokens > COMPACTION_TOKEN_THRESHOLD || last_assistant_message_stale?
+      end
+
+      def last_assistant_message_stale?
+        last_assistant_message_at = session_manager.last_assistant_message_at
+        last_assistant_message_at && Time.now - last_assistant_message_at > COMPACTION_IDLE_THRESHOLD_SECONDS
+      end
+
+      def drain_queue(queue)
+        session_manager.drain_message_queue(queue, mode: queue_drain_mode)
+      end
+
+      def emit(event, &block)
+        return unless block
+
+        block.call(event)
+      end
+    end
+  end
+end

data/lib/llm_gateway/agents/in_memory_session_manager.rb

@@ -0,0 +1,222 @@
+# frozen_string_literal: true
+
+require "securerandom"
+require "time"
+
+module LlmGateway
+  module Agents
+    class InMemorySessionManager
+      MESSAGE_QUEUED = :queued
+      MESSAGE_STARTED = :started
+      QUEUES = [ :steer, :follow_up, :next_turn ].freeze
+      DRAIN_MODES = [ :one_at_a_time, :all ].freeze
+
+      attr_reader :session_id, :session_start
+
+      def initialize(session_id = nil)
+        @state = :idle
+        @session_id = session_id
+        @message_queues = Hash.new { |hash, key| hash[key] = [] }
+      end
+
+      def busy!
+        @state = :busy
+      end
+
+      def idle!
+        @state = :idle
+      end
+
+      def drain_message_queue(queue = :next_turn, mode: :all)
+        messages = queued_messages(queue, mode)
+        messages.each { |message| push_message(message) }
+        messages
+      end
+
+      def queued_messages?(queue = :next_turn)
+        @message_queues[validate_queue!(queue)].any?
+      end
+
+      def push_message_to_queue(message, queue = :next_turn)
+        @message_queues[validate_queue!(queue)] << message
+      end
+
+      def busy?
+        @state == :busy
+      end
+
+      def validate_queue!(queue)
+        queue = queue.to_sym
+        raise ArgumentError, "Invalid queue mode: #{queue}" unless QUEUES.include?(queue)
+
+        queue
+      end
+
+      def validate_drain_mode!(mode)
+        mode = mode.to_sym
+        raise ArgumentError, "Invalid queue drain mode: #{mode}" unless DRAIN_MODES.include?(mode)
+
+        mode
+      end
+
+      def start_or_enqueue_user_message(payload, queue: :next_turn)
+        if busy?
+          push_message_to_queue(payload, queue)
+          MESSAGE_QUEUED
+        else
+          yield if block_given?
+          push_message(payload)
+          busy!
+          MESSAGE_STARTED
+        end
+      end
+
+      def push_message(payload)
+        payload = payload.deep_symbolize_keys
+
+        push_entry(
+          type: "message",
+          usage: message_usage(payload),
+          data: payload,
+        )
+      end
+
+      def push_entry(entry)
+        id = SecureRandom.uuid
+        new_entry = {
+          id: id,
+          parent_id: parent_id_for_new_entry,
+          timestamp: Time.now.iso8601,
+          **entry
+        }
+
+        persist_entry(new_entry)
+        new_entry
+      end
+
+      def active_messages
+        active_message_events.map { |event| event[:data] }
+      end
+
+      def last_message_id
+        message_events.last&.dig(:id)
+      end
+
+      def last_model_used
+        events.reverse.find { |event| event[:type] == "model_change" }&.dig(:model_id)
+      end
+
+      def last_reasoning_level_used
+        events.reverse.find { |event| event[:type] == "reasoning_change" }&.dig(:reasoning)
+      end
+
+      def events_until(event_id)
+        index = events.index { |event| event[:id] == event_id }
+        raise ArgumentError, "Event not found in session: #{event_id}" unless index
+
+        events[0..index]
+      end
+
+      def events
+        @events ||= [ new_session_event ]
+      end
+
+      def build_model_input_messages
+        return active_messages unless last_compaction_entry
+
+        [ last_compaction_entry[:data], *active_messages ]
+      end
+
+      def total_tokens
+        entry = active_message_events.reverse.find { |event| event.dig(:usage, :total_tokens) }
+        entry&.dig(:usage, :total_tokens) || 0
+      end
+
+      def last_assistant_message_at
+        entry = active_message_events.reverse.find { |event| event.dig(:data, :role) == "assistant" }
+        Time.parse(entry[:timestamp]) if entry
+      end
+
+      def compaction(adapter)
+        response = adapter.stream(
+          active_messages,
+          system: "Summarize the conversation so far for future context.",
+          tools: []
+        )
+        message = response.to_h
+
+        push_entry(
+          type: "compaction",
+          usage: message_usage(message),
+          data: message
+        )
+      end
+
+      private
+
+      def queued_messages(queue, mode)
+        queue = validate_queue!(queue)
+        case validate_drain_mode!(mode)
+        when :one_at_a_time
+          message = @message_queues[queue].shift
+          message ? [ message ] : []
+        when :all
+          @message_queues[queue].shift(@message_queues[queue].length)
+        end
+      end
+
+      def parent_id_for_new_entry
+        events.length.positive? ? events.last[:id] : nil
+      end
+
+      def message_events
+        events.select { |event| event[:type] == "message" }
+      end
+
+      def active_message_events
+        compaction_event = last_compaction_entry
+        return message_events unless compaction_event
+
+        compaction_index = events.index(compaction_event)
+        events[(compaction_index + 1)..].select { |event| event[:type] == "message" }
+      end
+
+      def last_compaction_entry
+        events.reverse.find { |event| event[:type] == "compaction" }
+      end
+
+      def message_usage(message)
+        usage = message[:usage] || message["usage"]
+        return {} unless usage
+
+        usage.transform_keys(&:to_sym)
+      end
+
+      def persist_entry(entry)
+        attributes = {
+          session_id: @session_id,
+          position: next_position,
+          id: entry[:id],
+          parent_id: entry[:parent_id],
+          timestamp: entry[:timestamp],
+          type: entry[:type],
+          usage: entry[:usage],
+          data: entry[:data]
+        }
+
+        events << entry
+        attributes
+      end
+
+      def next_position
+        events.length
+      end
+
+      def new_session_event
+        @session_id ||= SecureRandom.uuid
+        @session_start = Time.now.strftime("%Y%m%d_%H%M%S")
+        { type: "session", id: session_id, timestamp: session_start }
+      end
+    end
+  end
+end