rbs_goose 0.1.1 → 0.2.0

data/lib/rbs_goose/io/example_group.rb

@@ -3,22 +3,22 @@
 module RbsGoose
   module IO
     class ExampleGroup < Array
+      attr_accessor :error_messages
+
       class << self
         def load_from(base_path, code_dir: 'lib', sig_dir: 'sig', refined_dir: 'refined')
           new.tap do |group|
             Dir.glob('**/*.rb', base: ::File.join(base_path, code_dir)).each do |path|
-              group << Example.from_path(
-                ruby_path: ::File.join(code_dir, path),
-                rbs_path: to_rbs_path(path, sig_dir),
-                refined_rbs_dir: refined_dir,
-                base_path: base_path
-              )
+              group << load_example(base_path, code_dir, path, refined_dir, sig_dir)
             end
+            # Read error messages if exists
+            error_path = ::File.join(base_path, 'error_messages.txt')
+            group.error_messages = ::File.read(error_path).strip if ::File.exist?(error_path)
           end
         end
 
         def default_examples
-          example_dir = ::File.join(__dir__.to_s, '../examples')
+          example_dir = ::File.join(__dir__.to_s, '../../../examples')
           @default_examples ||= Dir.glob('*', base: example_dir).to_h do |dir|
             [dir.to_sym, load_from(::File.join(example_dir, dir))]
           end
@@ -26,6 +26,15 @@ module RbsGoose
 
         private
 
+        def load_example(base_path, code_dir, path, refined_dir, sig_dir)
+          Example.from_path(
+            ruby_path: ::File.join(code_dir, path),
+            rbs_path: to_rbs_path(path, sig_dir),
+            refined_rbs_dir: refined_dir,
+            base_path:
+          )
+        end
+
         def to_rbs_path(path, sig_dir)
           ::File.join(sig_dir, "#{path}s")
         end

data/lib/rbs_goose/io/target_group.rb

@@ -10,7 +10,7 @@ module RbsGoose
               group << TypedRuby.from_path(
                 ruby_path: ::File.join(code_dir, path),
                 rbs_path: ::File.join(sig_dir, "#{path}s"),
-                base_path: base_path
+                base_path:
               )
             end
           end

data/lib/rbs_goose/io/typed_ruby.rb

@@ -5,13 +5,13 @@ module RbsGoose
     class TypedRuby
       class << self
         def from_path(ruby_path:, rbs_path:, base_path:)
-          ruby = File.new(path: ruby_path, base_path: base_path)
+          ruby = File.new(path: ruby_path, base_path:)
           rbs = begin
-            File.new(path: rbs_path, base_path: base_path)
+            File.new(path: rbs_path, base_path:)
           rescue StandardError
             nil
           end
-          new(ruby: ruby, rbs: rbs)
+          new(ruby:, rbs:)
         end
       end
 

data/lib/rbs_goose/orthoses/fix_error.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require 'orthoses/outputable'
+
+module RbsGoose
+  module Orthoses
+    # Call RbsGoose::TypeInferrer to fix steep type errors.
+    #   use TbsGoose::Orthoses::FixError, code_dir: 'lib', sig_dir: 'sig' do |config|
+    #     config.use_open_ai('open_ai_access_token')
+    #   end
+    class FixError
+      def initialize(loader, code_dir: 'lib', sig_dir: 'sig', &)
+        @loader = loader
+        @code_dir = code_dir
+        @sig_dir = sig_dir
+        RbsGoose.configure(&) if block_given?
+      end
+
+      def call
+        @loader.call.tap do
+          RbsGoose.fix_error(code_dir: @code_dir, sig_dir: @sig_dir)
+        end
+      end
+    end
+  end
+end

data/lib/rbs_goose/orthoses/infer.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require 'orthoses/outputable'
+
+module RbsGoose
+  module Orthoses
+    # Call RbsGoose::TypeInferrer to infer RBS type signatures.
+    #   use TbsGoose::Orthoses::Infer, code_dir: 'lib', sig_dir: 'sig' do |config|
+    #     config.use_open_ai('open_ai_access_token')
+    #   end
+    class Infer
+      def initialize(loader, code_dir: 'lib', sig_dir: 'sig', &)
+        @loader = loader
+        @code_dir = code_dir
+        @sig_dir = sig_dir
+        RbsGoose.configure(&) if block_given?
+      end
+
+      def call
+        @loader.call.tap do
+          RbsGoose.run(code_dir: @code_dir, sig_dir: @sig_dir)
+        end
+      end
+    end
+  end
+end

data/lib/rbs_goose/orthoses.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+require_relative 'orthoses/fix_error'
+require_relative 'orthoses/infer'
+
+module RbsGoose
+  module Orthoses
+  end
+end

data/lib/rbs_goose/templates/base.rb

@@ -0,0 +1,77 @@
+# frozen_string_literal: true
+
+require 'langchain'
+
+module RbsGoose
+  module Templates
+    class Base
+      def initialize(instruction:, example_groups:)
+        @instruction = instruction
+        @example_groups = example_groups
+      end
+
+      def format(**args)
+        Langchain::Prompt::FewShotPromptTemplate.new(
+          prefix:,
+          suffix:,
+          example_prompt:,
+          examples: example_groups.map { transform_example_group(_1) },
+          input_variables:
+        ).format(**format_args(args)).strip
+      end
+
+      def format_system_prompt
+        prefix.strip
+      end
+
+      def format_user_prompt(**args)
+        Langchain::Prompt::FewShotPromptTemplate.new(
+          prefix: '',
+          suffix:,
+          example_prompt:,
+          examples: example_groups.map { transform_example_group(_1) },
+          input_variables:
+        ).format(**format_args(args)).strip
+      end
+
+      def parse_result(result)
+        result.scan(/```.+?```/m).map { IO::File.from_markdown(_1) }
+      end
+
+      private
+
+      def input_template_string
+        raise NotImplementedError
+      end
+
+      def input_variables
+        raise NotImplementedError
+      end
+
+      def transform_example_group(example_group)
+        raise NotImplementedError
+      end
+
+      def format_args(args)
+        raise NotImplementedError
+      end
+
+      def prefix
+        instruction
+      end
+
+      def suffix
+        "#{input_template_string}\n"
+      end
+
+      def example_prompt
+        Langchain::Prompt::PromptTemplate.new(
+          template: "#{input_template_string}\n{refined_rbs_list}",
+          input_variables: input_variables + ['refined_rbs_list']
+        )
+      end
+
+      attr_reader :instruction, :example_groups
+    end
+  end
+end

data/lib/rbs_goose/templates/fix_error_template.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require 'langchain'
+
+module RbsGoose
+  module Templates
+    class FixErrorTemplate < Base
+      private
+
+      def input_variables
+        %w[typed_ruby_list error_messages]
+      end
+
+      def input_template_string
+        <<~TEMPLATE.strip
+          ========Input========
+          {typed_ruby_list}
+
+          ========Errors========
+          {error_messages}
+
+
+          ========Output========
+        TEMPLATE
+      end
+
+      def format_args(args)
+        {
+          typed_ruby_list: args[:typed_ruby_list].join("\n"),
+          error_messages: args[:error_messages]
+        }
+      end
+
+      def transform_example_group(example_group)
+        {
+          typed_ruby_list: example_group.map(&:typed_ruby).join("\n"),
+          error_messages: example_group.error_messages,
+          refined_rbs_list: example_group.map(&:refined_rbs).join("\n")
+        }
+      end
+    end
+  end
+end

data/lib/rbs_goose/templates/infer_template.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require 'langchain'
+
+module RbsGoose
+  module Templates
+    class InferTemplate < Base
+      private
+
+      def input_variables
+        %w[typed_ruby_list]
+      end
+
+      def input_template_string
+        "========Input========\n{typed_ruby_list}\n\n========Output========"
+      end
+
+      def format_args(args)
+        { typed_ruby_list: args[:typed_ruby_list].join("\n") }
+      end
+
+      def transform_example_group(example_group)
+        {
+          typed_ruby_list: example_group.map(&:typed_ruby).join("\n"),
+          refined_rbs_list: example_group.map(&:refined_rbs).join("\n")
+        }
+      end
+    end
+  end
+end

data/lib/rbs_goose/templates.rb

@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
-require_relative 'templates/default_template'
+require_relative 'templates/base'
+require_relative 'templates/fix_error_template'
+require_relative 'templates/infer_template'
 
 module RbsGoose
   module Templates

data/lib/rbs_goose/type_inferrer.rb

@@ -1,13 +1,84 @@
 # frozen_string_literal: true
 
+require 'steep'
+require 'steep/cli'
+require 'benchmark'
+
 require_relative 'templates'
 
 module RbsGoose
   class TypeInferrer
     def infer(target_group)
-      template = RbsGoose.configuration.template
-      result = RbsGoose.llm.complete(prompt: template.format(target_group)).completion
+      call_llm(RbsGoose.infer_template, { typed_ruby_list: target_group })
+    end
+
+    def fix_error(target_group)
+      error_messages = steep_check
+      call_llm(RbsGoose.fix_error_template, { typed_ruby_list: target_group, error_messages: })
+    end
+
+    private
+
+    def call_llm(template, format_args)
+      result =
+        case RbsGoose.llm_mode.to_sym
+        when :complete
+          call_llm_complete(format_args, template)
+        when :chat
+          call_llm_chat(format_args, template)
+        else
+          raise "Invalid LLM mode: #{RbsGoose.llm_mode}"
+        end
       template.parse_result(result)
     end
+
+    def call_llm_complete(format_args, template)
+      prompt = template.format(**format_args)
+      llm_debug(prompt) do
+        RbsGoose.llm.complete(prompt:)
+      end.completion
+    end
+
+    def call_llm_chat(format_args, template)
+      messages = [
+        { role: 'user', content: template.format(**format_args) }
+      ]
+      llm_debug(messages.map { "role: #{_1[:role]}\ncontent:\n#{_1[:content]}" }.join("\n")) do
+        RbsGoose.llm.chat(messages:)
+      end.chat_completion
+    end
+
+    def llm_debug(prompt)
+      return yield if ENV['DEBUG'].nil?
+
+      puts "!!!!!!!! Prompt !!!!!!!!\n\n#{prompt}\n\n"
+      result = nil
+      sec = Benchmark.realtime do
+        result = yield
+      end
+      puts "!!!!!!!! Stats !!!!!!!!\n\n  spend: #{sec}[s]\n  prompt_tokens: #{result.prompt_tokens}\n  completion_tokens: #{result.completion_tokens}\n" # rubocop:disable Layout/LineLength
+      puts "!!!!!!!! Result !!!!!!!!\n\n#{result.chat_completion}\n\n"
+      result
+    end
+
+    def steep_check
+      stdin, stdout, stderr = io_stubs
+      disable_rainbow do
+        steep_cli = Steep::CLI.new(stdout:, stdin:, stderr:, argv: [])
+        steep_cli.process_check
+      end
+      stdout.string
+    end
+
+    def io_stubs
+      [StringIO.new(''), StringIO.new(+'', 'w+'), StringIO.new(+'', 'w+')]
+    end
+
+    def disable_rainbow
+      rainbow_enabled = Rainbow.enabled
+      Rainbow.enabled = false
+      yield
+      Rainbow.enabled = rainbow_enabled
+    end
   end
 end

data/lib/rbs_goose/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module RbsGoose
-  VERSION = '0.1.1'
+  VERSION = '0.2.0'
 end

data/lib/rbs_goose.rb

@@ -2,6 +2,7 @@
 
 require_relative 'rbs_goose/configuration'
 require_relative 'rbs_goose/io'
+require_relative 'rbs_goose/orthoses'
 require_relative 'rbs_goose/type_inferrer'
 require_relative 'rbs_goose/version'
 
@@ -13,8 +14,8 @@ module RbsGoose
   class << self
     extend Forwardable
 
-    def configure(&block)
-      @configuration = Configuration.new(&block)
+    def configure(&)
+      @configuration = Configuration.new(&)
     end
 
     def reset_configuration
@@ -23,7 +24,7 @@ module RbsGoose
 
     def run(code_dir: 'lib', sig_dir: 'sig', base_path: ::Dir.pwd)
       puts "Run RbsGoose.(Code Directory: #{code_dir}, Signature Directory: #{sig_dir})"
-      target_group = RbsGoose::IO::TargetGroup.load_from(base_path, code_dir: code_dir, sig_dir: sig_dir)
+      target_group = RbsGoose::IO::TargetGroup.load_from(base_path, code_dir:, sig_dir:)
       RbsGoose::TypeInferrer.new.infer(target_group).each do |refined_rbs|
         puts "write refined rbs to #{refined_rbs.path}\n"
         refined_rbs.write
@@ -31,8 +32,27 @@ module RbsGoose
       end
     end
 
+    def fix_error(code_dir: 'lib', sig_dir: 'sig', base_path: ::Dir.pwd)
+      puts "Run FixError.(Code Directory: #{code_dir}, Signature Directory: #{sig_dir})"
+      target_group = RbsGoose::IO::TargetGroup.load_from(base_path, code_dir:, sig_dir:)
+      RbsGoose::TypeInferrer.new.fix_error(target_group).each do |refined_rbs|
+        puts "write refined rbs to #{refined_rbs.path}\n"
+        refined_rbs.write
+        puts "done.\n\n"
+      end
+    end
+
+    def infer_template
+      configuration.infer_template.build_template
+    end
+
+    def fix_error_template
+      configuration.fix_error_template.build_template
+    end
+
     attr_reader :configuration
 
-    def_delegators :configuration, :llm, :instruction, :example_groups
+    def_delegators :configuration, :infer_instruction, :infer_example_groups, :llm_mode
+    def_delegator :configuration, :llm_client, :llm
   end
 end

data/rbs_collection.lock.yaml

@@ -1,12 +1,16 @@
 ---
 path: ".gem_rbs_collection"
 gems:
+- name: abbrev
+  version: '0'
+  source:
+    type: stdlib
 - name: activesupport
   version: '7.0'
   source:
     type: git
     name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
+    revision: 7a105f52053ce1c708b605dfa9c1ab8473424036
     remote: https://github.com/ruby/gem_rbs_collection.git
     repo_dir: gems
 - name: addressable
@@ -14,7 +18,7 @@ gems:
   source:
     type: git
     name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
+    revision: 7a105f52053ce1c708b605dfa9c1ab8473424036
     remote: https://github.com/ruby/gem_rbs_collection.git
     repo_dir: gems
 - name: ast
@@ -22,7 +26,7 @@ gems:
   source:
     type: git
     name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
+    revision: 7a105f52053ce1c708b605dfa9c1ab8473424036
     remote: https://github.com/ruby/gem_rbs_collection.git
     repo_dir: gems
 - name: baran
@@ -42,7 +46,7 @@ gems:
   source:
     type: git
     name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
+    revision: 7a105f52053ce1c708b605dfa9c1ab8473424036
     remote: https://github.com/ruby/gem_rbs_collection.git
     repo_dir: gems
 - name: connection_pool
@@ -50,9 +54,13 @@ gems:
   source:
     type: git
     name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
+    revision: 7a105f52053ce1c708b605dfa9c1ab8473424036
     remote: https://github.com/ruby/gem_rbs_collection.git
     repo_dir: gems
+- name: csv
+  version: '0'
+  source:
+    type: stdlib
 - name: date
   version: '0'
   source:
@@ -66,7 +74,7 @@ gems:
   source:
     type: git
     name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
+    revision: 7a105f52053ce1c708b605dfa9c1ab8473424036
     remote: https://github.com/ruby/gem_rbs_collection.git
     repo_dir: gems
 - name: fileutils
@@ -82,13 +90,21 @@ gems:
   source:
     type: git
     name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
+    revision: 7a105f52053ce1c708b605dfa9c1ab8473424036
     remote: https://github.com/ruby/gem_rbs_collection.git
     repo_dir: gems
 - name: json
   version: '0'
   source:
     type: stdlib
+- name: listen
+  version: '3.2'
+  source:
+    type: git
+    name: ruby/gem_rbs_collection
+    revision: 7a105f52053ce1c708b605dfa9c1ab8473424036
+    remote: https://github.com/ruby/gem_rbs_collection.git
+    repo_dir: gems
 - name: logger
   version: '0'
   source:
@@ -105,62 +121,50 @@ gems:
   version: '0'
   source:
     type: stdlib
-- name: parallel
-  version: '1.20'
+- name: net-http
+  version: '0'
   source:
-    type: git
-    name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
-    remote: https://github.com/ruby/gem_rbs_collection.git
-    repo_dir: gems
+    type: stdlib
+- name: net-protocol
+  version: '0'
+  source:
+    type: stdlib
+- name: optparse
+  version: '0'
+  source:
+    type: stdlib
+- name: orthoses
+  version: 1.13.0
+  source:
+    type: rubygems
 - name: parser
   version: '3.2'
   source:
     type: git
     name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
+    revision: 7a105f52053ce1c708b605dfa9c1ab8473424036
     remote: https://github.com/ruby/gem_rbs_collection.git
     repo_dir: gems
+- name: pathname
+  version: '0'
+  source:
+    type: stdlib
 - name: rainbow
   version: '3.0'
   source:
     type: git
     name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
+    revision: 7a105f52053ce1c708b605dfa9c1ab8473424036
     remote: https://github.com/ruby/gem_rbs_collection.git
     repo_dir: gems
-- name: rake
-  version: '13.0'
+- name: rbs
+  version: 3.4.1
   source:
-    type: git
-    name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
-    remote: https://github.com/ruby/gem_rbs_collection.git
-    repo_dir: gems
-- name: regexp_parser
-  version: '2.8'
-  source:
-    type: git
-    name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
-    remote: https://github.com/ruby/gem_rbs_collection.git
-    repo_dir: gems
-- name: rubocop
-  version: '1.57'
-  source:
-    type: git
-    name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
-    remote: https://github.com/ruby/gem_rbs_collection.git
-    repo_dir: gems
-- name: rubocop-ast
-  version: '1.30'
+    type: rubygems
+- name: rdoc
+  version: '0'
   source:
-    type: git
-    name: ruby/gem_rbs_collection
-    revision: ec140aff951e3af39846eed4525ba5d040d294fa
-    remote: https://github.com/ruby/gem_rbs_collection.git
-    repo_dir: gems
+    type: stdlib
 - name: securerandom
   version: '0'
   source:
@@ -169,8 +173,16 @@ gems:
   version: '0'
   source:
     type: stdlib
+- name: steep
+  version: 1.6.0
+  source:
+    type: rubygems
+- name: strscan
+  version: '0'
+  source:
+    type: stdlib
 - name: tiktoken_ruby
-  version: 0.0.6
+  version: 0.0.8
   source:
     type: rubygems
 - name: time
@@ -181,6 +193,10 @@ gems:
   version: '0'
   source:
     type: stdlib
+- name: tsort
+  version: '0'
+  source:
+    type: stdlib
 - name: uri
   version: '0'
   source:

data/renovate.json

@@ -10,6 +10,11 @@
     {
       "groupName": "rubocop",
       "matchPackagePatterns": [ "rubocop" ]
+    },
+    {
+      "matchUpdateTypes": ["minor", "patch"],
+      "matchCurrentVersion": "!/^0/",
+      "automerge": true
     }
   ]
 }