llm_gateway 0.6.0 → 0.8.0

data/lib/llm_gateway/adapters/openai/responses/stream_mapper.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "json"
+
 require_relative "../../stream_mapper"
 
 module LlmGateway
@@ -23,10 +25,16 @@ module LlmGateway
               response_created_patches(data[:response])
             when "response.output_item.added"
               output_item_added_patches(data)
+            when "response.output_item.done"
+              output_item_done_patches(data)
             when "response.content_part.added"
               content_part_added_patches(data)
-            when "response.content_part.done"
+            when "response.content_part.done", "response.output_text.done"
               content_part_done_patches(data)
+            when "response.code_interpreter_call_code.delta"
+              code_interpreter_code_delta_patches(data)
+            when "response.code_interpreter_call.in_progress", "response.code_interpreter_call.interpreting", "response.code_interpreter_call.completed", "response.code_interpreter_call_code.done"
+              []
             when "response.output_text.delta"
               [ { type: :text_delta, delta: data[:delta] || "" } ]
             when "response.function_call_arguments.delta"
@@ -84,6 +92,38 @@ module LlmGateway
                   name: item[:name]
                 }
               ]
+            when "code_interpreter_call"
+              state = code_interpreter_state[data[:output_index] || 0] = {
+                id: item[:id],
+                container_id: item[:container_id],
+                outputs: item[:outputs],
+                input_opened: false,
+                input_closed: false
+              }
+              container_id_to_tool_id[state[:container_id]] = state[:id] if state[:container_id]
+
+              [
+                {
+                  type: :tool_start,
+                  delta: "",
+                  id: item[:id],
+                  name: "code_interpreter_call",
+                  tool_type: "server_tool_use"
+                }
+              ]
+            else
+              []
+            end
+          end
+
+          def output_item_done_patches(data)
+            item = data[:item] || {}
+
+            case item[:type]
+            when "code_interpreter_call"
+              code_interpreter_done_patches(data[:output_index] || 0, item)
+            when "message"
+              container_file_citation_patches(item)
             else
               []
             end
@@ -100,7 +140,83 @@ module LlmGateway
             part = data[:part] || {}
             return [] unless part.empty? || part[:type] == "output_text"
 
-            [ { type: :text_end, delta: "" } ]
+            citations = container_file_citation_patches(data)
+            return citations unless accumulator.active_block_type == :text
+
+            [ { type: :text_end, delta: "" } ] + citations
+          end
+
+          def code_interpreter_code_delta_patches(data)
+            output_index = data[:output_index] || 0
+            state = code_interpreter_state[output_index] ||= {
+              id: nil,
+              container_id: nil,
+              outputs: nil,
+              input_opened: false,
+              input_closed: false
+            }
+            delta = escape_json_string_fragment(data[:delta] || "")
+            delta = "{\"code\":\"#{delta}" unless state[:input_opened]
+            state[:input_opened] = true
+
+            [ { type: :tool_delta, delta: } ]
+          end
+
+          def code_interpreter_done_patches(output_index, item)
+            state = code_interpreter_state[output_index] ||= {}
+            state[:id] ||= item[:id]
+            state[:container_id] = item[:container_id] if item.key?(:container_id)
+            state[:outputs] = item[:outputs] if item.key?(:outputs)
+            container_id_to_tool_id[state[:container_id]] = state[:id] if state[:container_id] && state[:id]
+            return [] if state[:input_closed]
+
+            opening = state[:input_opened] ? "" : "{\"code\":\""
+            state[:input_opened] = true
+            closing = "\"," + JSON.generate(container_id: state[:container_id], outputs: state[:outputs])[1..]
+            state[:input_closed] = true
+
+            [
+              { type: :tool_delta, delta: opening + closing },
+              { type: :tool_end, delta: "" }
+            ]
+          end
+
+          def container_file_citation_patches(data)
+            extract_annotations(data).filter_map do |annotation|
+              next unless annotation[:type] == "container_file_citation"
+
+              container_id = annotation[:container_id]
+              file_id = annotation[:file_id]
+              filename = annotation[:filename]
+              tool_id = container_id_to_tool_id[container_id]
+              next unless tool_id
+
+              key = [ tool_id, container_id, file_id, filename ]
+              next if emitted_citation_keys[key]
+
+              emitted_citation_keys[key] = true
+              {
+                type: :tool_result_start,
+                delta: JSON.generate(container_id:, file_id:, filename:),
+                tool_use_id: tool_id,
+                name: "container_file_citation_tool_result"
+              }
+            end.flat_map { |start| [ start, { type: :tool_result_end, delta: "" } ] }
+          end
+
+          def extract_annotations(data)
+            annotations = []
+            annotations.concat(Array(data[:annotations]))
+            annotations.concat(Array(data.dig(:part, :annotations)))
+            annotations.concat(Array(data.dig(:item, :annotations)))
+            Array(data.dig(:item, :content)).each do |content_part|
+              annotations.concat(Array(content_part[:annotations])) if content_part.is_a?(Hash)
+            end
+            annotations
+          end
+
+          def escape_json_string_fragment(value)
+            JSON.generate(value)[1...-1]
           end
 
           def response_completed_patches(response)
@@ -162,7 +278,19 @@ module LlmGateway
           end
 
           def tool_seen?
-            accumulator.blocks.any? { |content_block| content_block && content_block[:type] == "tool_use" }
+            accumulator.blocks.any? { |content_block| content_block && [ "tool_use", "server_tool_use" ].include?(content_block[:type]) }
+          end
+
+          def code_interpreter_state
+            @code_interpreter_state ||= {}
+          end
+
+          def container_id_to_tool_id
+            @container_id_to_tool_id ||= {}
+          end
+
+          def emitted_citation_keys
+            @emitted_citation_keys ||= {}
           end
         end
       end

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