roast-ai 0.4.9 → 0.4.10

data/lib/roast/dsl/cog/input.rb

@@ -0,0 +1,30 @@
+# typed: true
+# frozen_string_literal: true
+
+module Roast
+  module DSL
+    class Cog
+      # Abstract parent class for inputs provided to a cog when it runs.
+      # Cogs should extend this class with their own input types
+      # Input classes should be instantiatable with a no-arg constructor,
+      # and expose methods to incrementally set their values.
+      class Input
+        class InputError < Roast::Error; end
+        class InvalidInputError < InputError; end
+
+        # Validate that the input instance has all required parameters set in an acceptable manner.
+        # Inheriting cog must implement this for its input class.
+        #: () -> void
+        def validate!
+          raise NotImplementedError
+        end
+
+        # Use the value returned from the cog's input block to attempt to coerce the input to a valid state.
+        # to coerce the input to a valid state.
+        # Inheriting cog may implement this for its input class; it is optional
+        #: (untyped) -> void
+        def coerce(input_return_value); end
+      end
+    end
+  end
+end

data/lib/roast/dsl/cog/output.rb

@@ -0,0 +1,24 @@
+# typed: true
+# frozen_string_literal: true
+
+module Roast
+  module DSL
+    class Cog
+      # Generic output from running a cog.
+      # Cogs should extend this class with their own output types.
+      class Output
+        #: bool
+        attr_reader :success
+
+        #: bool
+        attr_reader :abort
+
+        #: (?success: bool, ?abort_workflow: bool) -> void
+        def initialize(success: true, abort_workflow: false)
+          @success = success #: bool
+          @abort = abort_workflow #: bool
+        end
+      end
+    end
+  end
+end

data/lib/roast/dsl/cog/registry.rb

@@ -0,0 +1,39 @@
+# typed: true
+# frozen_string_literal: true
+
+module Roast
+  module DSL
+    class Cog
+      class Registry
+        class CogRegistryError < Roast::Error; end
+        class CouldNotDeriveCogNameError < CogRegistryError; end
+
+        def initialize
+          @cogs = {}
+          create_registration(Cogs::Cmd)
+          create_registration(Cogs::Chat)
+          create_registration(Cogs::Execute)
+        end
+
+        #: Hash[Symbol, singleton(Cog)]
+        attr_reader :cogs
+
+        #: (singleton(Roast::DSL::Cog)) -> void
+        def use(cog_class)
+          reg = create_registration(cog_class)
+          cogs[reg.first] = reg.second
+        end
+
+        private
+
+        #: (singleton(Roast::DSL::Cog)) -> Array(Symbol, singleton(Cog))
+        def create_registration(cog_class)
+          cog_class_name = cog_class.name
+          raise CouldNotDeriveCogNameError if cog_class_name.nil?
+
+          [cog_class_name.demodulize.underscore.to_sym, cog_class]
+        end
+      end
+    end
+  end
+end

data/lib/roast/dsl/cog/stack.rb

@@ -7,8 +7,9 @@ module Roast
       class Stack
         delegate :map, :push, :size, :empty?, to: :@queue
 
+        #: () -> void
         def initialize
-          @queue = []
+          @queue = [] #: Array[Cog]
         end
 
         #: () -> Roast::DSL::Cog?

data/lib/roast/dsl/cog/store.rb

@@ -9,17 +9,20 @@ module Roast
 
         delegate :[], to: :store
 
+        #: Hash[Symbol, Cog]
+        attr_reader :store
+
+        #: () -> void
+        def initialize
+          @store = {}
+        end
+
         #: (Symbol, Roast::DSL::Cog) -> Roast::DSL::Cog
         def insert(id, inst)
           raise CogAlreadyDefinedError if store.key?(id)
 
           store[id] = inst
         end
-
-        #: () -> Hash[Symbol, Roast::DSL::Cog]
-        def store
-          @store ||= {}
-        end
       end
     end
   end

data/lib/roast/dsl/cog.rb

@@ -4,67 +4,88 @@
 module Roast
   module DSL
     class Cog
-      class CogAlreadyRanError < StandardError; end
+      class CogError < Roast::Error; end
 
-      class << self
-        def on_create
-          eigen = self
-          proc do |instance_name = Random.uuid, &action|
-            #: self as Roast::DSL::WorkflowExecutionContext
-            add_cog_instance(instance_name, eigen.new(instance_name, action))
-          end
-        end
-
-        def on_config
-          eigen = self
-          proc do |cog_name = nil, &configuration|
-            #: self as Roast::DSL::ConfigContext
-            config_object = if cog_name.nil?
-              fetch_execution_scope(eigen)
-            else
-              fetch_or_create_cog_config(eigen, cog_name)
-            end
-
-            config_object.instance_exec(&configuration) if configuration
-          end
-        end
+      class CogAlreadyRanError < CogError; end
 
+      class << self
+        #: () -> singleton(Cog::Config)
         def config_class
           @config_class ||= find_child_config_or_default
         end
 
+        #: () -> singleton(Cog::Input)
+        def input_class
+          @input_class ||= find_child_input_or_default #: singleton(Cog::Input)?
+        end
+
         private
 
+        #: () -> singleton(Cog::Config)
         def find_child_config_or_default
           config_constant = "#{name}::Config"
           const_defined?(config_constant) ? const_get(config_constant) : Cog::Config # rubocop:disable Sorbet/ConstantsFromStrings
         end
+
+        #: () -> singleton(Cog::Input)
+        def find_child_input_or_default
+          input_constant = "#{name}::Input"
+          const_defined?(input_constant) ? const_get(input_constant) : Cog::Input # rubocop:disable Sorbet/ConstantsFromStrings
+        end
       end
 
-      attr_reader :name, :output
+      #: Symbol
+      attr_reader :name
 
+      #: Cog::Output?
+      attr_reader :output
+
+      #: (Symbol, ^(Cog::Input) -> untyped) -> void
       def initialize(name, cog_input_proc)
         @name = name
-        @cog_input_proc = cog_input_proc
-        @finished = false
+        @cog_input_proc = cog_input_proc #: ^(Cog::Input) -> untyped
+        @output = nil #: Cog::Output?
+        @finished = false #: bool
+
+        # Make sure a config is always defined, so we don't have to worry about nils
+        @config = self.class.config_class.new #: Cog::Config
       end
 
-      def run!(config, cog_execution_context)
+      #: (Cog::Config, CogInputContext) -> void
+      def run!(config, input_context)
         raise CogAlreadyRanError if ran?
 
         @config = config
-        @output = execute(cog_execution_context.instance_exec(&@cog_input_proc))
+        input_instance = self.class.input_class.new
+        input_return = input_context.instance_exec(input_instance, &@cog_input_proc) if @cog_input_proc
+        coerce_and_validate_input!(input_instance, input_return)
+        @output = execute(input_instance)
         @finished = true
       end
 
+      #: () -> bool
       def ran?
         @finished
       end
 
       # Inheriting cog must implement this
+      #: (Cog::Input) -> Cog::Output
       def execute(input)
         raise NotImplementedError
       end
+
+      private
+
+      #: (Cog::Input, untyped) -> void
+      def coerce_and_validate_input!(input, return_value)
+        # Check if the input is already valid
+        input.validate!
+      rescue Cog::Input::InvalidInputError
+        # If it's not valid, attempt to coerce if possible
+        input.coerce(return_value)
+        # Re-validate because coerce! should not be responsible for validation
+        input.validate!
+      end
     end
   end
 end

data/lib/roast/dsl/cog_input_context.rb

@@ -0,0 +1,9 @@
+# typed: true
+# frozen_string_literal: true
+
+module Roast
+  module DSL
+    # Context in which the individual cog input blocks within the `execute` block of a workflow definition are evaluated
+    class CogInputContext; end
+  end
+end

data/lib/roast/dsl/cog_input_manager.rb

@@ -0,0 +1,47 @@
+# typed: true
+# frozen_string_literal: true
+
+module Roast
+  module DSL
+    # Context in which an individual cog block within the `execute` block of a workflow is evaluated
+    class CogInputManager
+      class CogOutputAccessError < Roast::Error; end
+      class CogNotYetRunError < CogOutputAccessError; end
+      class CogNotDefinedError < CogOutputAccessError; end
+
+      #: (Cog::Registry, Cog::Store) -> void
+      def initialize(cog_registry, cogs)
+        @cog_registry = cog_registry
+        @cogs = cogs
+        @context = CogInputContext.new
+        bind_registered_cogs
+      end
+
+      #: CogInputContext
+      attr_reader :context
+
+      private
+
+      #: () -> void
+      def bind_registered_cogs
+        @cog_registry.cogs.keys.each(&method(:bind_cog))
+      end
+
+      #: (Symbol) -> void
+      def bind_cog(cog_method_name)
+        cog_output_method = method(:cog_output)
+        @context.instance_eval do
+          define_singleton_method(cog_method_name, proc { |cog_name| cog_output_method.call(cog_name) })
+        end
+      end
+
+      #: (Symbol) -> Cog::Output
+      def cog_output(cog_name)
+        @cogs[cog_name].tap do |cog|
+          raise CogNotYetRunError unless cog.ran?
+          raise CogNotDefinedError unless cog.output.present?
+        end.output
+      end
+    end
+  end
+end

data/lib/roast/dsl/cogs/chat.rb

@@ -0,0 +1,78 @@
+# typed: true
+# frozen_string_literal: true
+
+module Roast
+  module DSL
+    module Cogs
+      class Chat < Cog
+        class Input < Cog::Input
+          #: String?
+          attr_accessor :prompt
+
+          #: () -> void
+          def initialize
+            super
+            @prompt = nil #: String?
+          end
+
+          #: () -> void
+          def validate!
+            raise Cog::Input::InvalidInputError, "'prompt' is required" unless prompt.present?
+          end
+
+          #: (untyped) -> void
+          def coerce!(input_return_value)
+            if input_return_value.is_a?(String)
+              self.prompt = input_return_value
+            end
+          end
+        end
+
+        class Output < Cog::Output
+          #: String
+          attr_reader :response
+
+          #: (String response) -> void
+          def initialize(response)
+            super()
+            @response = response
+          end
+        end
+
+        class Config < Cog::Config
+          #: () -> String?
+          def openai_api_key
+            @values[:openai_api_key] ||= ENV["OPENAI_API_KEY"]
+          end
+
+          #: () -> String?
+          def openai_api_base_url
+            @values[:openai_api_base_url] ||= ENV["OPENAI_API_BASE_URL"]
+          end
+        end
+
+        #: (Input) -> Output
+        def execute(input)
+          config = @config #: as Config
+          RubyLLM.configure do |ruby_llm_config|
+            ruby_llm_config.openai_api_key = config.openai_api_key
+            ruby_llm_config.openai_api_base = config.openai_api_base_url
+          end
+
+          chat = RubyLLM.chat
+          resp = chat.ask(input.prompt)
+          puts "Model: #{resp.model_id}"
+          puts "Role: #{resp.role}"
+          puts "Input Tokens: #{resp.input_tokens}"
+          puts "Output Tokens: #{resp.output_tokens}"
+
+          chat.messages.each do |message|
+            puts "[#{message.role.to_s.upcase}] #{message.content}"
+          end
+
+          Output.new(resp.content)
+        end
+      end
+    end
+  end
+end

data/lib/roast/dsl/cogs/cmd.rb

@@ -5,49 +5,126 @@ module Roast
   module DSL
     module Cogs
       class Cmd < Cog
-        class Output
+        class Input < Cog::Input
           #: String?
-          attr_reader :command_output
+          attr_accessor :command
 
-          #: String?
+          #: Array[String]
+          attr_accessor :args
+
+          #: () -> void
+          def initialize
+            super
+            @command = nil
+            @args = []
+          end
+
+          #: () -> void
+          def validate!
+            raise Cog::Input::InvalidInputError, "'command' is required" unless command.present?
+          end
+
+          #: (String | Array[untyped]) -> void
+          def coerce(input_return_value)
+            case input_return_value
+            when String
+              self.command = input_return_value
+            when Array
+              input_return_value.map!(&:to_s)
+              self.command = input_return_value.shift
+              self.args = input_return_value
+            end
+          end
+        end
+
+        class Output < Cog::Output
+          #: String
+          attr_reader :out
+
+          #: String
           attr_reader :err
 
-          #: Process::Status?
+          #: Process::Status
           attr_reader :status
 
-          #: (
-          #|  String? output,
-          #|  String? error,
-          #|  Process::Status? status
-          #| ) -> void
-          def initialize(output, error, status)
-            @command_output = output
-            @err = error
-            @status = status
+          #: ( String, String, Process::Status) -> void
+          def initialize(out, err, status)
+            super()
+            @out = out #: String
+            @err = err #: String
+            @status = status #: Process::Status
           end
         end
 
         class Config < Cog::Config
           #: () -> void
           def print_all!
-            @values[:print_all] = true
+            @values[:print_stdout] = true
+            @values[:print_stderr] = true
+          end
+
+          #: () -> void
+          def print_stdout!
+            @values[:print_stdout] = true
+          end
+
+          #: () -> void
+          def print_stderr!
+            @values[:print_stderr] = true
+          end
+
+          #: () -> bool
+          def print_stdout?
+            !!@values[:print_stdout]
           end
 
           #: () -> bool
-          def print_all?
-            !!@values[:print_all]
+          def print_stderr?
+            !!@values[:print_stderr]
           end
 
+          #: () -> void
           def display!
             print_all!
           end
         end
 
-        #: (String) -> Output
+        #: (Input) -> Output
         def execute(input)
-          result = Output.new(*Roast::Helpers::CmdRunner.capture3(input))
-          puts result.command_output if @config.print_all?
-          result
+          config = @config #: as Config
+          result = T.unsafe(Roast::Helpers::CmdRunner).popen3(input.command, *input.args) do |stdin, stdout, stderr, wait_thread|
+            stdin.close
+            command_output = ""
+            command_error = ""
+
+            # Thread to read and accumulate stdout
+            stdout_thread = Thread.new do
+              stdout.each_line do |line|
+                command_output += line
+                $stdout.puts(line) if config.print_stdout?
+              end
+            rescue IOError => e
+              Roast::Helpers::Logger.debug("IOError while reading stdout: #{e.message}")
+            end
+
+            # Thread to read and accumulate stderr
+            stderr_thread = Thread.new do
+              stderr.each_line do |line|
+                command_error += line
+                $stderr.puts(line) if config.print_stderr?
+              end
+            rescue IOError => e
+              Roast::Helpers::Logger.debug("IOError while reading stderr: #{e.message}")
+            end
+
+            # Wait for threads to finish
+            stdout_thread.join
+            stderr_thread.join
+
+            [command_output, command_error, wait_thread.value]
+          end #: as [String, String, Process::Status]
+
+          Output.new(*result)
         end
       end
     end

data/lib/roast/dsl/cogs/execute.rb

@@ -0,0 +1,46 @@
+# typed: true
+# frozen_string_literal: true
+
+module Roast
+  module DSL
+    module Cogs
+      class Execute < Cog
+        class Input < Cog::Input
+          #: Symbol?
+          attr_accessor :scope
+
+          #: () -> void
+          def initialize
+            super
+            @scope = nil
+          end
+
+          #: () -> void
+          def validate!
+            raise Cog::Input::InvalidInputError, "'scope' is required" unless scope.present?
+          end
+
+          #: (Symbol) -> void
+          def coerce(input_return_value)
+            self.scope = input_return_value
+          end
+        end
+
+        #: (Symbol, ^(Input) -> untyped, ^(Input) -> void) -> void
+        def initialize(name, cog_input_proc, trigger)
+          # NOTE: Sorbet expects the proc passed to super to be declared as taking a Cog::Input explicitly,
+          # not a subclass of Cog::Input.
+          cog_input_proc = cog_input_proc #: as ^(Cog::Input) -> untyped
+          super(name, cog_input_proc)
+          @trigger = trigger
+        end
+
+        #: (Input) -> Cog::Output
+        def execute(input)
+          @trigger.call(input)
+          Cog::Output.new
+        end
+      end
+    end
+  end
+end

data/lib/roast/dsl/cogs/graph.rb

@@ -8,12 +8,12 @@ module Roast
         #: Proc
         attr_reader :block
 
-        #: (Symbol) { (Roast::DSL::Cogs::Graph) -> void } -> void
+        #: (Symbol) { (Cog::Input) -> untyped } -> void
         def initialize(name, &block)
           @name = name
           @block = block
           @graph = Roast::Graph.new
-          super(name, proc {})
+          super(name, block)
         end
 
         #: () -> void
@@ -36,7 +36,7 @@ module Roast
         end
 
         # Populates the provided graph in-place, with the definition of how to populate in the block
-        #: (Roast::DSL::Cogs::Graph) -> void
+        #: (Cog::Input) -> void
         def populate!(graph)
           return if @block.nil?
 

data/lib/roast/dsl/config_context.rb

@@ -3,52 +3,7 @@
 
 module Roast
   module DSL
-    class ConfigContext
-      def initialize(cogs, config_proc)
-        @cogs = cogs
-        @executor_scoped_configs = {}
-        @cog_scoped_configs = {}
-        @config_proc = config_proc
-      end
-
-      def fetch_merged_config(cog_class, name = nil)
-        # All configs have an entry, even if it's empty.
-        configs = fetch_execution_scope(cog_class)
-        instance_configs = fetch_cog_config(cog_class, name) unless name.nil?
-        configs = configs.merge(instance_configs) if instance_configs
-        configs
-      end
-
-      def prepare!
-        bind_default_cogs
-        instance_eval(&@config_proc)
-      end
-
-      #: () -> void
-      def bind_default_cogs
-        bind_cog(Cogs::Cmd, :cmd)
-      end
-
-      def fetch_cog_config(cog_class, name)
-        @cog_scoped_configs[cog_class][name]
-      end
-
-      def fetch_or_create_cog_config(cog_class, name)
-        @cog_scoped_configs[cog_class][name] = cog_class.config_class.new unless @cog_scoped_configs.key?(name)
-        @cog_scoped_configs[cog_class][name]
-      end
-
-      def fetch_execution_scope(cog_class)
-        @executor_scoped_configs[cog_class]
-      end
-
-      def bind_cog(cog_class, method_name)
-        @cog_scoped_configs[cog_class] = {}
-        @executor_scoped_configs[cog_class] = cog_class.config_class.new
-        instance_eval do
-          define_singleton_method(method_name, &cog_class.on_config)
-        end
-      end
-    end
+    # Context in which the `config` blocks of a workflow definition are evaluated
+    class ConfigContext; end
   end
 end