llm.rb 12.1.0 → 12.3.0

data/lib/llm/repl/bar.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+class LLM::Repl
+  ##
+  # The {LLM::Repl::Bar LLM::Repl::Bar} class renders a
+  # small progress bar for the REPL. It is used to show
+  # the remaining size of the model's context window in
+  # a compact form near the input line.
+  # @api private
+  class Bar
+    ##
+    # @return [String]
+    OCCUPIED = "█"
+
+    ##
+    # @return [String]
+    FREE = " "
+
+    ##
+    # @param [Integer] used
+    # @param [Integer] total
+    # @param [Integer] width
+    # @return [LLM::Repl::Bar]
+    def initialize(used:, total:, width: 10)
+      @width = width
+      @label, @filled = remainder(used, total)
+    end
+
+    ##
+    # @return [String]
+    def to_s
+      bar = "#{OCCUPIED * filled}#{FREE * (width - filled)}"
+      "│#{bar}│ #{label}"
+    end
+
+    private
+
+    ##
+    # @param [Integer] used
+    # @param [Integer] total
+    # @return [[String, Integer]]
+    def remainder(used, total)
+      return ["???", width] if total <= 0
+      diff = total - used
+      return ["0%", 0] if diff <= 0
+      remaining = (diff.fdiv(total) * 100).round(2)
+      ["#{remaining}%", ((remaining / 100) * width).round]
+    end
+
+    attr_reader :label, :filled, :width
+  end
+end

data/lib/llm/repl/input.rb

@@ -6,59 +6,172 @@ class LLM::Repl
   # the editable input line shown at the bottom of the REPL.
   # @api private
   class Input
+    CTRL_A    = 1
+    CTRL_E    = 5
     UP        = Curses::Key::UP
     DOWN      = Curses::Key::DOWN
+    LEFT      = Curses::Key::LEFT
+    RIGHT     = Curses::Key::RIGHT
     ENTER     = [Curses::Key::ENTER, 10, 13]
     BACKSPACE = [Curses::Key::BACKSPACE, 127]
-    EOF       = [nil, 4]
+    EOF       = [4]
 
     ##
-    # @param [String, Symbol] provider
+    # @param [LLM::Agent] agent
     # @return [LLM::Repl::Input]
-    def initialize(provider)
-      @provider = provider
+    def initialize(agent, options = {})
+      @agent = agent
+      @provider = agent.llm.name
       @buffer = +""
+      @cursor = 0
+      @scroll = 0
+      @height = options.fetch(:height, 3)
     end
 
     ##
     # @param [LLM::Repl::Window] window
-    # @return [String, nil]
-    def readline(window)
-      catch(:done) do
-        @buffer.clear
-        loop do
-          on_char(window, window.getch)
-          window.redraw
-        end
-      end
-    end
-
-    ##
-    # @return [String]
-    def to_s
-      "> #{@buffer}"
-    end
-
-    private
-
+    # @param [Object] char
+    # @return [Symbol, nil]
     def on_char(window, char)
       if EOF.include?(char)
-        throw(:done, nil)
+        :exit
       elsif BACKSPACE.include?(char)
-        @buffer.chop!
+        backspace
+        :backspace
       elsif ENTER.include?(char)
-        buf = @buffer.dup
-        @buffer.clear
-        throw(:done, buf)
+        :submit
       elsif char == UP
         window.scroll_up
+        :up
       elsif char == DOWN
         window.scroll_down
+        :down
+      elsif char == CTRL_A
+        move_start
+        :ctrl_a
+      elsif char == CTRL_E
+        move_end
+        :ctrl_e
+      elsif char == LEFT
+        move_left
+        :left
+      elsif char == RIGHT
+        move_right
+        :right
       elsif String === char
-        @buffer << char
+        insert(char)
+        :char
       else
-        # ???
+        nil
       end
     end
+
+    ##
+    # @return [String]
+    def to_s
+      "#{@provider}> #{@buffer}"
+    end
+
+    ##
+    # @return [Integer]
+    def cursor
+      prompt.length + @cursor
+    end
+
+    ##
+    # @return [Integer]
+    def height
+      @height
+    end
+
+    ##
+    # Returns the visible lines of the input buffer,
+    # wrapped at the given column width. The viewport
+    # follows the cursor so the cursor line is always
+    # visible.
+    # @param [Integer] cols
+    # @return [Array<String>]
+    def lines(cols)
+      sync_scroll(cols)
+      text = to_s
+      chunks = text.chars.each_slice(cols).map(&:join)
+      chunks = [""] if chunks.empty?
+      chunks[@scroll, height] || []
+    end
+
+    ##
+    # Returns the cursor position as [line, column] within
+    # the visible viewport.
+    # @param [Integer] cols
+    # @return [Array(Integer, Integer)]
+    def cursor_pos(cols)
+      sync_scroll(cols)
+      [(cursor / cols) - @scroll, cursor % cols]
+    end
+
+    ##
+    # @return [void]
+    def move_start
+      @cursor = 0
+    end
+
+    ##
+    # @return [void]
+    def move_end
+      @cursor = [0, @buffer.size].max
+    end
+
+    ##
+    # @return [void]
+    def move_left
+      @cursor = [@cursor - 1, 0].max
+    end
+
+    ##
+    # @return [void]
+    def move_right
+      @cursor = [@cursor + 1, @buffer.size].min
+    end
+
+    ##
+    # @return [String]
+    def take
+      @buffer.dup.tap do
+        @buffer.clear
+        @cursor = 0
+        @scroll = 0
+      end
+    end
+
+    private
+
+    ##
+    # Adjusts @scroll so the cursor line is visible within
+    # the viewport.
+    def sync_scroll(cols)
+      total_lines = [1, (to_s.length.to_f / cols).ceil].max
+      cursor_line = cursor / cols
+      if cursor_line < @scroll
+        @scroll = cursor_line
+      elsif cursor_line >= (@scroll + height)
+        @scroll = (cursor_line - height) + 1
+      end
+      @scroll = [[@scroll, (total_lines - height)].min, 0].max
+    end
+
+    def prompt
+      "#{@provider}> "
+    end
+
+    def insert(char)
+      @buffer.insert(@cursor, char)
+      @cursor += char.length
+    end
+
+    def backspace
+      return if @cursor <= 0
+      @buffer.slice!(@cursor - 1)
+      @cursor -= 1
+    end
   end
 end

data/lib/llm/repl/markdown.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+class LLM::Repl
+  ##
+  # This class is designed to represent a markdown
+  # string (typically from a model's response) as a
+  # tree of objects where each object contains a piece
+  # of text, and also optional style information for
+  # that text (eg bold, underscore, ...)
+  class Markdown
+    ##
+    # @param [String] text
+    # @return [LLM::Repl::Markdown]
+    def initialize(text)
+      @doc = Kramdown::Document.new(text)
+      @ast = []
+    end
+
+    ##
+    # @return [Array<Hash>]
+    def ast
+      @ast.tap do
+        ##
+        # Recurisvely travels the markdown document and
+        # populates the `@ast` variable along the way.
+        # The AST is composed of structured data that
+        # carries both text and styling information that
+        # is applied by the UI thread.
+        walk(@doc.root)
+
+        ##
+        # This is required because the AST collects
+        # empty nodes towards the end of the tree.
+        # If we don't pop them we end up with excessive
+        # amount of newlines between turns.
+        last = @ast.last
+        while last and last[:text].to_s.strip.empty?
+          @ast.pop
+          last = @ast.last
+        end
+      end
+    end
+
+    private
+
+    ##
+    # Recursively walk from the head node to the
+    # tail node. This method mutates the `@ast`
+    # variable. A future refactor might be worthwhile
+    # since this method is implemented with side effects,
+    # but it probably could return the ast instead.
+    def walk(node, attrs = nil)
+      case node.type
+      when :root
+        node.children.each { walk(_1, attrs) }
+      when :text
+        emit(node.value.to_s, attrs)
+      when :p
+        node.children.each { walk(_1, attrs) }
+        emit("\n\n", attrs)
+      when :header
+        emit("\n", attrs)
+        node.children.each { walk(_1, Curses::A_BOLD) }
+        emit("\n", attrs)
+      when :strong
+        node.children.each { walk(_1, Curses::A_BOLD) }
+      when :em
+        node.children.each { walk(_1, Curses::A_UNDERLINE) }
+      when :codespan
+        emit(node.value, Curses::A_REVERSE)
+      when :codeblock
+        emit(node.value, Curses::A_REVERSE)
+        emit("\n\n", attrs)
+      when :br
+        emit("\n", attrs)
+      else
+        node.children.each { walk(_1, attrs) }
+      end
+    end
+
+    def emit(text, attrs)
+      @ast.push({text: text.to_s, attrs:}.compact)
+    end
+  end
+end

data/lib/llm/repl/status.rb

@@ -9,18 +9,29 @@ class LLM::Repl
     ##
     # @param [String, Symbol] provider
     # @return [LLM::Repl::Status]
-    def initialize(provider)
-      @provider = provider
+    def initialize(agent)
+      @agent = agent
+      @provider = agent.llm.name
       @text = "idle"
     end
 
     ##
     # @return [String]
-    attr_reader :provider
+    def context_bar
+      LLM::Repl::Bar.new(
+        used: @agent.usage.total_tokens,
+        total: @agent.context_window
+      ).to_s
+    end
+
+    ##
+    # @return [String]
+    def cost
+      "$#{@agent.cost}"
+    end
 
     ##
-    # @param [String] value
-    # @return [void]
+    # @return [String]
     attr_accessor :text
 
     ##

data/lib/llm/repl/stream.rb

@@ -11,8 +11,10 @@ class LLM::Repl
     ##
     # @param [LLM::Repl] repl
     # @return [LLM::Repl::Stream]
-    def initialize(repl)
+    def initialize(repl, queue)
       @repl = repl
+      @_queue = queue
+      @buffer = +""
     end
 
     ##
@@ -20,7 +22,8 @@ class LLM::Repl
     #  One or more chars
     # @return [void]
     def on_content(chars)
-      @repl.write(chars)
+      @buffer << chars
+      @_queue.push [:stream, @buffer]
     end
 
     ##
@@ -29,9 +32,9 @@ class LLM::Repl
     # @return [void]
     def on_tool_call(tool, error)
       if error
-        @repl.status = "tool error: #{tool.name}"
+        @_queue.push [:status, "tool not found: #{tool.name}"]
       else
-        @repl.status = "tool: #{tool.name}"
+        @_queue.push [:status, "tool: #{tool.name}"]
       end
     end
 
@@ -40,7 +43,14 @@ class LLM::Repl
     # @param [LLM::Function::Return] result
     # @return [void]
     def on_tool_return(_tool, result)
-      @repl.status = "tool done: #{result.name}"
+      @_queue.push [:status, "tool done: #{result.name}"]
+    end
+
+    ##
+    # Empty the accumulated buffer
+    # @return [void]
+    def empty!
+      @buffer.clear
     end
   end
 end

data/lib/llm/repl/transcript.rb

@@ -2,30 +2,70 @@
 
 class LLM::Repl
   ##
-  # The {LLM::Repl::Transcript LLM::Repl::Transcript} class
-  # stores streamed output for the REPL.
-  # @api private
+  # This class maintains conversation state that includes
+  # the conversation itself, and metadata associated with
+  # the conversation.
+  #
+  # Internally it maintains an array where each element
+  # represents a row, and each element in a row is a Hash
+  # that describes a piece of text and any styles that might
+  # be applied to it by the UI thread.
+  #
+  # It also maintains a cursor that tracks the active row
+  # by its index number. The streaming path reuses a single
+  # row by overwriting its contents repeatedly.
   class Transcript
     WIDTH = 80
 
     ##
     # @return [LLM::Repl::Transcript]
     def initialize
-      @lines = [+""]
+      @rows = [[]]
+      @cursor = nil
+      @snapshot = nil
       @offset = 0
     end
 
     ##
     # @param [String] chars
+    # @param [Object] attrs
+    # @param [Symbol] method
     # @return [void]
-    def write(chars)
-      chars.each_char { write_char(_1) }
+    def write(chars, attrs = nil, method: :append)
+      chunks = [{text: chars.to_s, attrs:}.compact]
+      self.method(method).call(chunks)
+    end
+
+    ##
+    # Appends Markdown to the transcript.
+    # @param [String] chars
+    # @param [Symbol] method
+    # @return [void]
+    def markdown(chars, method: :append)
+      chunks = LLM::Repl::Markdown.new(chars).ast
+      self.method(method).call(chunks)
+    end
+
+    ##
+    # Start the transcript.
+    # @return [void]
+    def start
+      @cursor = @rows.size - 1
+      @snapshot = @rows.map(&:dup)
+    end
+
+    ##
+    # Finish the transcript.
+    # @return [void]
+    def finish
+      @cursor = nil
+      @snapshot = nil
     end
 
     ##
     # @return [void]
     def scroll_up(height)
-      max = [@lines.size - height, 0].max
+      max = [rows.size - height, 0].max
       @offset = [@offset + 1, max].min
     end
 
@@ -39,23 +79,72 @@ class LLM::Repl
     # @param [Integer] height
     # @return [Array<String>]
     def visible(height)
-      last = @lines.size - 1 - @offset
+      all = rows
+      last = all.size - 1 - @offset
       first = [last - height + 1, 0].max
-      @lines[first..last] || []
+      all[first..last] || []
     end
 
     private
 
-    def write_char(char)
-      if char == "\n"
-        @offset += 1 if @offset > 0
-        @lines << +""
-      elsif char == " " and @lines.last.length >= WIDTH
-        @offset += 1 if @offset > 0
-        @lines << +""
-      else
-        @lines.last << char
+    ##
+    # Appends a new row
+    # @param [Array<{text: String, attrs?: Integer}>] chunks
+    #  One or more chunks.
+    # @return [void]
+    def append(chunks)
+      chunks.each { wrap(_1, @rows) }
+    end
+
+    ##
+    # Replaces the content of the active row
+    # @param [Array<{text: String, attrs?: Integer}>] chunks
+    #  One or more chunks.
+    # @return [void]
+    def replace(chunks)
+      @rows = @snapshot.map(&:dup)
+      append(chunks)
+    end
+
+    ##
+    # Given a chunk this method wraps text at
+    # around 80 columns: a new row starts when
+    # the current character is " " and the sum
+    # of all characters in that row is greater
+    # than 80 columns.
+    def wrap(chunk, rows)
+      attrs = chunk[:attrs]
+      chunk[:text].to_s.each_char do |char|
+        if char == "\n"
+          rows << []
+        elsif char == " " and sum(rows.last) >= WIDTH
+          rows << []
+        else
+          rows.last << {text: char, attrs:}.compact
+        end
       end
     end
+
+    ##
+    # @api private
+    def rows
+      @rows.dup.tap do |rows|
+        ##
+        # Discard empty rows that would otherwise
+        # be rendered as newlines by the UI thread.
+        # It's not the most elegant way to deal with
+        # this and we probably shouldn't allow it to
+        # happen in the first place.
+        while rows.size > 1 and rows.last.empty?
+          rows.pop
+        end
+      end
+    end
+
+    ##
+    # @api private
+    def sum(row)
+      row.sum { _1[:text].to_s.length }
+    end
   end
 end

data/lib/llm/repl/window.rb

@@ -37,6 +37,7 @@ class LLM::Repl
       Curses.cbreak
       Curses.noecho
       Curses.stdscr.keypad(true)
+      Curses.stdscr.nodelay = true
       yield
     ensure
       Curses.close_screen
@@ -46,8 +47,9 @@ class LLM::Repl
     # @return [void]
     def redraw
       Curses.clear
-      draw_status
-      draw_transcript
+      draw_divider(offset: 5)
+      draw_status(offset: input.height + 1)
+      draw_transcript(offset: 0)
       draw_input
       Curses.refresh
     end
@@ -55,7 +57,7 @@ class LLM::Repl
     ##
     # @return [Integer]
     def rows
-      [Curses.lines - 3, 1].max
+      [Curses.lines - (input.height + 4), 1].max
     end
 
     ##
@@ -78,28 +80,52 @@ class LLM::Repl
 
     private
 
-    def draw_status
-      Curses.setpos(0, 0)
+    def draw_status(offset:)
+      Curses.setpos(Curses.lines - offset, 0)
+      Curses.clrtoeol
       Curses.addstr(status.to_s)
-      provider = status.provider.to_s
-      Curses.setpos(0, [columns - provider.length, 0].max)
-      Curses.addstr(provider)
+      context = status.context_bar
+      Curses.setpos(Curses.lines - offset, [(columns - context.length) / 2, 0].max)
+      Curses.addstr(context)
+      cost = status.cost.to_s
+      Curses.setpos(Curses.lines - offset, [columns - cost.length, 0].max)
+      Curses.addstr(cost)
     end
 
-    def draw_transcript
-      visible = transcript.visible(rows)
-      visible.each_with_index do |row, index|
-        Curses.setpos(index + 2, 0)
-        Curses.addstr(row)
-      end
+    def draw_divider(offset:)
+      Curses.setpos(Curses.lines - offset, 0)
+      Curses.clrtoeol
+      Curses.addstr("─" * Curses.cols)
     end
 
     def draw_input
-      Curses.setpos(Curses.lines - 1, 0)
-      Curses.clrtoeol
-      Curses.addstr(input.to_s)
+      cols = columns
+      rows = input.lines(cols)
+      rows.each.with_index do |line, idx|
+        Curses.setpos((Curses.lines - input.height) + idx, 0)
+        Curses.clrtoeol
+        Curses.addstr(line)
+      end
+      line, col = input.cursor_pos(cols)
+      Curses.setpos((Curses.lines - input.height) + line, col)
     end
 
+    def draw_transcript(offset:)
+      rows = transcript.visible(self.rows)
+      rows.each.with_index(offset) do |row, index|
+        Curses.setpos(index, 0)
+        Curses.clrtoeol
+        row.each do |chunk|
+          text, attrs = chunk.values_at(:text, :attrs)
+          Curses.attron(attrs) if attrs
+          Curses.addstr(text)
+          Curses.attroff(attrs) if attrs
+        end
+      end
+    end
+
+    ##
+    # @return [Integer]
     def columns
       Curses.cols
     end