llm_gateway 0.5.0 → 0.7.0

data/lib/llm_gateway/adapters/structs.rb

@@ -9,35 +9,6 @@ class BaseStruct < Dry::Struct
   transform_keys(&:to_sym)
 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)
-
-  attribute :type, EventType
-  attribute :delta, Types::Coercible::String.default { "" }
-  attribute :content_index, Types::Integer
-end
-
-
-class AssistantToolStartEvent < AssistantStreamEvent
-  attribute :id, Types::String
-  attribute :name, Types::String
-  attribute :content_index, Types::Integer
-end
-
-
-class AssistantStreamReasoningEvent < AssistantStreamEvent
-  attribute :signature, Types::Coercible::String.default { "" }
-  attribute :content_index, Types::Integer
-end
-
-class AssistantStreamMessageEvent < BaseStruct
-  EventType = Types::Coercible::Symbol.enum(:message_start, :message_delta, :message_end)
-
-  attribute :type, EventType
-  attribute :delta, Types::Coercible::Hash.default { {} }
-  attribute :usage_increment, Types::Coercible::Hash.default { {} }
-end
-
 class TextContent < BaseStruct
   attribute :type, Types::String.enum("text")
   attribute :text, Types::String
@@ -81,18 +52,142 @@ 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 AssistantMessage < BaseStruct
+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
+  attribute? :role, Types::String.enum("assistant").optional
+  attribute :timestamp, Types::Integer
+  attribute? :stop_reason, Types::String.enum("stop", "length", "tool_use", "toolUse", "error", "aborted").optional
+  attribute? :content, Types::Array.of(ContentBlock).optional
+
+  def self.new(attributes = {})
+    attrs = attributes.to_h.transform_keys(&:to_sym)
+    attrs[:content] = Array(attrs[:content]).map { |block| build_content_block(block) } if attrs.key?(:content)
+    super(attrs)
+  end
+
+  def self.build_content_block(block)
+    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"
+      TextContent.new(block)
+    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"])
+    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
+
+  private_class_method :build_content_block
+end
+
+class AssistantStreamEvent < BaseStruct
+  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 { "" }
+  attribute :content_index, Types::Integer
+  attribute :partial, Types.Instance(PartialAssistantMessage)
+
+  def content
+    case type
+    when :text_end
+      finalized_content_block&.text
+    when :reasoning_end
+      finalized_content_block&.reasoning
+    when :tool_end
+      finalized_content_block
+    end
+  end
+
+  def text
+    content if type == :text_end
+  end
+
+  def reasoning
+    content if type == :reasoning_end
+  end
+
+  def tool_call
+    finalized_content_block if type == :tool_end
+  end
 
+  alias tool tool_call
+
+  private
+
+  def finalized_content_block
+    partial.content&.[](content_index)
+  end
+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
+
+class AssistantStreamReasoningEvent < AssistantStreamEvent
+  attribute :signature, Types::Coercible::String.default { "" }
+  attribute :content_index, Types::Integer
+end
+
+class AssistantMessage < PartialAssistantMessage
   attribute :id, Types::String
   attribute :model, Types::String
   attribute :usage, Types::Hash
@@ -103,12 +198,6 @@ class AssistantMessage < BaseStruct
   attribute? :error_message, Types::String.optional
   attribute :content, Types::Array.of(ContentBlock)
 
-  def self.new(attributes)
-    attrs = attributes.to_h.transform_keys(&:to_sym)
-    attrs[:content] = Array(attrs[:content]).map { |block| build_content_block(block) }
-    super(attrs)
-  end
-
   def to_h
     result = {
       id: id,
@@ -120,26 +209,22 @@ class AssistantMessage < BaseStruct
       api: api,
       content: content.map(&:to_h)
     }
+    result[:timestamp] = timestamp unless timestamp.nil?
     result[:error_message] = error_message unless error_message.nil?
     result
   end
+end
 
-  def self.build_content_block(block)
-    return block if block.is_a?(TextContent) || block.is_a?(ReasoningContent) || block.is_a?(ToolCall)
+class AssistantStreamMessageEvent < BaseStruct
+  EventType = Types::Coercible::Symbol.enum(:message_start, :message_delta)
 
-    case block[:type] || block["type"]
-    when "text"
-      TextContent.new(block)
-    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"])
-    when "tool_use"
-      ToolCall.new(block)
-    else
-      raise ArgumentError, "Unsupported content block type: #{block[:type] || block['type']}"
-    end
-  end
+  attribute :type, EventType
+  attribute :delta, Types::Coercible::Hash.default { {} }
+  attribute :usage, Types::Coercible::Hash.default { {} }
+  attribute :partial, Types.Instance(PartialAssistantMessage)
+end
 
-  private_class_method :build_content_block
+class AssistantStreamMessageEndEvent < BaseStruct
+  attribute :type, Types::Coercible::Symbol.enum(:message_end)
+  attribute :message, Types.Instance(AssistantMessage)
 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