roast-ai 0.4.9 → 0.5.0

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

@@ -0,0 +1,71 @@
+# typed: true
+# frozen_string_literal: true
+
+module Roast
+  module DSL
+    class Cog
+      # Registry for managing available cogs in a workflow
+      class Registry
+        # Parent class for all cog registry errors
+        class CogRegistryError < Roast::Error; end
+
+        # Raised when a cog name cannot be derived from a cog class
+        class CouldNotDeriveCogNameError < CogRegistryError; end
+
+        # Initialize a new cog registry with standard cogs
+        #
+        # Automatically registers all system cogs and standard cogs provided by core Roast.
+        #
+        # #### See Also
+        # - `use`
+        # - `cogs`
+        #
+        def initialize
+          @cogs = {}
+          use(SystemCogs::Call)
+          use(SystemCogs::Map)
+          use(SystemCogs::Repeat)
+          use(Cogs::Cmd)
+          use(Cogs::Chat)
+          use(Cogs::Agent)
+          use(Cogs::Ruby)
+        end
+
+        # Hash mapping cog names to cog classes
+        #
+        # #### See Also
+        # - `use`
+        #
+        #: Hash[Symbol, singleton(Cog)]
+        attr_reader :cogs
+
+        # Register a cog class for use in workflows
+        #
+        # Adds the provided cog class to the registry, deriving its name from the class name.
+        # The cog name is the underscored, demodulized version of the class name
+        # (e.g., `Roast::DSL::Cogs::MyCustomCog` becomes `:my_custom_cog`).
+        #
+        # Raises `CouldNotDeriveCogNameError` if the cog class name cannot be determined.
+        #
+        # #### See Also
+        # - `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

@@ -5,10 +5,11 @@ module Roast
   module DSL
     class Cog
       class Stack
-        delegate :map, :push, :size, :empty?, to: :@queue
+        delegate :each, :empty?, :last, :map, :push, :size, to: :@queue
 
+        #: () -> void
         def initialize
-          @queue = []
+          @queue = [] #: Array[Cog]
         end
 
         #: () -> Roast::DSL::Cog?

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

@@ -7,18 +7,21 @@ module Roast
       class Store
         class CogAlreadyDefinedError < Roast::Error; end
 
-        delegate :[], to: :store
+        delegate :[], :key?, to: :store
 
-        #: (Symbol, Roast::DSL::Cog) -> Roast::DSL::Cog
-        def insert(id, inst)
-          raise CogAlreadyDefinedError if store.key?(id)
+        #: Hash[Symbol, Cog]
+        attr_reader :store
 
-          store[id] = inst
+        #: () -> void
+        def initialize
+          @store = {}
         end
 
-        #: () -> Hash[Symbol, Roast::DSL::Cog]
-        def store
-          @store ||= {}
+        #: (Cog) -> Roast::DSL::Cog
+        def insert(cog)
+          raise CogAlreadyDefinedError, cog.name if store.key?(cog.name)
+
+          store[cog.name] = cog
         end
       end
     end

data/lib/roast/dsl/cog.rb

@@ -4,67 +4,144 @@
 module Roast
   module DSL
     class Cog
-      class CogAlreadyRanError < StandardError; end
+      class CogError < Roast::Error; end
+
+      class CogAlreadyStartedError < CogError; 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
+        #: () -> singleton(Cog::Config)
+        def config_class
+          @config_class ||= find_child_config_or_default
         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
+        #: () -> singleton(Cog::Input)
+        def input_class
+          @input_class ||= find_child_input_or_default #: singleton(Cog::Input)?
         end
 
-        def config_class
-          @config_class ||= find_child_config_or_default
+        #: () -> Symbol
+        def generate_fallback_name
+          Random.uuid.to_sym
         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?
+        @skipped = false #: bool
+        @failed = false #: bool
+
+        # Make sure a config is always defined, so we don't have to worry about nils
+        @config = self.class.config_class.new #: untyped
+      end
+
+      #: (Async::Barrier, Cog::Config, CogInputContext, untyped, Integer) -> Async::Task
+      def run!(barrier, config, input_context, executor_scope_value, executor_scope_index)
+        raise CogAlreadyStartedError if @task
+
+        @task = barrier.async(finished: false) do |task|
+          task.annotate("#{self.class.name.not_nil!.demodulize.camelcase} Cog: #{@name}")
+          @config = config
+          input_instance = self.class.input_class.new
+          input_return = input_context.instance_exec(
+            input_instance, executor_scope_value, executor_scope_index, &@cog_input_proc
+          ) if @cog_input_proc
+          coerce_and_validate_input!(input_instance, input_return)
+          @output = execute(input_instance)
+        rescue ControlFlow::SkipCog
+          # TODO: do something with the message passed to skip!
+          @skipped = true
+        rescue ControlFlow::FailCog => e
+          # TODO: do something with the message passed to fail!
+          @failed = true
+          # TODO: better / cleaner handling in the workflow execution manager for a workflow failure
+          #   just re-raising this exception for now
+          raise e if config.abort_on_failure?
+        rescue ControlFlow::Next, ControlFlow::Break => e
+          @skipped = true
+          raise e
+        rescue StandardError => e
+          @failed = true
+          raise e
+        end
+      end
+
+      #: () -> void
+      def wait
+        @task&.wait
+      rescue
+        # Do nothing if the cog's task raised an exception. That is handled elsewhere.
+      end
+
+      #: () -> bool
+      def started?
+        @task.present?
+      end
+
+      #: () -> bool
+      def skipped?
+        @skipped
       end
 
-      def run!(config, cog_execution_context)
-        raise CogAlreadyRanError if ran?
+      #: () -> bool
+      def failed?
+        @failed || !!@task&.failed?
+      end
 
-        @config = config
-        @output = execute(cog_execution_context.instance_exec(&@cog_input_proc))
-        @finished = true
+      #: () -> bool
+      def stopped?
+        !!@task&.stopped?
       end
 
-      def ran?
-        @finished
+      #: () -> bool
+      def succeeded?
+        # NOTE: explicitly checking `@output == nil` because the `ruby` cog's Output class delegates
+        # all missing methods to its `value`, which may be nil when the Output object is actually present.
+        @output != nil && @task&.finished? # rubocop:disable Style/NonNilCheck
       end
 
+      protected
+
       # 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,44 @@
+# 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
+      include SystemCogs::Call::InputContext
+      include SystemCogs::Map::InputContext
+
+      class CogInputContextError < Roast::Error; end
+
+      class ContextNotFoundError < CogInputContextError; end
+
+      #: (?String?) -> void
+      def skip!(message = nil)
+        raise ControlFlow::SkipCog, message
+      end
+
+      #: (?String?) -> void
+      def fail!(message = nil)
+        raise ControlFlow::FailCog, message
+      end
+
+      #: (?String?) -> void
+      def next!(message = nil)
+        raise ControlFlow::Next, message
+      end
+
+      #: (?String?) -> void
+      def break!(message = nil)
+        raise ControlFlow::Break, message
+      end
+
+      #: (String, ?Hash) -> String
+      def template(path, args = {})
+        path = "prompts/#{path}.md.erb" unless File.exist?(path)
+        fail!("The prompt #{path} could not be found") unless File.exist?(path)
+
+        ERB.new(File.read(path)).result_with_hash(args)
+      end
+    end
+  end
+end

data/lib/roast/dsl/cog_input_manager.rb

@@ -0,0 +1,156 @@
+# 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 CogDoesNotExistError < CogOutputAccessError; end
+
+      class CogNotYetRunError < CogOutputAccessError; end
+
+      class CogSkippedError < CogOutputAccessError; end
+
+      class CogFailedError < CogOutputAccessError; end
+
+      class CogStoppedError < CogOutputAccessError; end
+
+      #: (Cog::Registry, Cog::Store, WorkflowContext) -> void
+      def initialize(cog_registry, cogs, workflow_context)
+        @cog_registry = cog_registry
+        @cogs = cogs
+        @workflow_context = workflow_context
+        @context = CogInputContext.new
+        bind_registered_cogs
+        bind_workflow_context
+      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_question_method_name = (cog_method_name.to_s + "?").to_sym
+        cog_bang_method_name = (cog_method_name.to_s + "!").to_sym
+        cog_output_method = method(:cog_output)
+        cog_output_question_method = method(:cog_output?)
+        cog_output_bang_method = method(:cog_output!)
+        @context.instance_eval do
+          define_singleton_method(cog_method_name, proc { |cog_name| cog_output_method.call(cog_name) })
+          define_singleton_method(cog_question_method_name, proc { |cog_name| cog_output_question_method.call(cog_name) })
+          define_singleton_method(cog_bang_method_name, proc { |cog_name| cog_output_bang_method.call(cog_name) })
+        end
+      end
+
+      #: (Symbol) -> Cog::Output?
+      def cog_output(cog_name)
+        cog_output!(cog_name)
+      rescue CogOutputAccessError => e
+        # Even this method should raise an exception if the requested cog does not exist at all
+        raise e if e.is_a?(CogDoesNotExistError)
+
+        nil
+      end
+
+      #: (Symbol) -> bool
+      def cog_output?(cog_name)
+        !cog_output(cog_name).nil?
+      end
+
+      #: (Symbol) -> Cog::Output
+      def cog_output!(cog_name)
+        raise CogDoesNotExistError, cog_name unless @cogs.key?(cog_name)
+
+        @cogs[cog_name].tap do |cog|
+          cog.wait # attempting to access the output of a running cog will block until that cog completes
+          raise CogSkippedError, cog_name if cog.skipped?
+          raise CogFailedError, cog_name if cog.failed?
+          raise CogStoppedError, cog_name if cog.stopped?
+          raise CogNotYetRunError, cog_name unless cog.succeeded?
+        end.output.deep_dup
+      end
+
+      #: () -> void
+      def bind_workflow_context
+        target_bang_method = method(:target!)
+        targets_method = method(:targets)
+        arg_question_method = method(:arg?)
+        args_method = method(:args)
+        kwarg_method = method(:kwarg)
+        kwarg_bang_method = method(:kwarg!)
+        kwarg_question_method = method(:kwarg?)
+        kwargs_method = method(:kwargs)
+        tmpdir_method = method(:tmpdir)
+        @context.instance_eval do
+          define_singleton_method(:target!, proc { target_bang_method.call })
+          define_singleton_method(:targets, proc { targets_method.call })
+          define_singleton_method(:arg?, proc { |value| arg_question_method.call(value) })
+          define_singleton_method(:args, proc { args_method.call })
+          define_singleton_method(:kwarg, proc { |key| kwarg_method.call(key) })
+          define_singleton_method(:kwarg!, proc { |key| kwarg_bang_method.call(key) })
+          define_singleton_method(:kwarg?, proc { |key| kwarg_question_method.call(key) })
+          define_singleton_method(:kwargs, proc { kwargs_method.call })
+          define_singleton_method(:tmpdir, proc { tmpdir_method.call })
+        end
+      end
+
+      #: () -> String
+      def target!
+        raise ArgumentError, "expected exactly one target" unless @workflow_context.params.targets.length == 1
+
+        @workflow_context.params.targets.first #: as String
+      end
+
+      #: () -> Array[String]
+      def targets
+        @workflow_context.params.targets.dup
+      end
+
+      #: (Symbol) -> bool
+      def arg?(value)
+        @workflow_context.params.args.include?(value)
+      end
+
+      #: () -> Array[Symbol]
+      def args
+        @workflow_context.params.args.dup
+      end
+
+      #: (Symbol) -> String?
+      def kwarg(key)
+        @workflow_context.params.kwargs[key]
+      end
+
+      #: (Symbol) -> String
+      def kwarg!(key)
+        raise ArgumentError, "expected keyword argument '#{key}' to be present" unless @workflow_context.params.kwargs.include?(key)
+
+        @workflow_context.params.kwargs[key] #: as String
+      end
+
+      #: (Symbol) -> bool
+      def kwarg?(key)
+        @workflow_context.params.kwargs.include?(key)
+      end
+
+      #: () -> Hash[Symbol, String]
+      def kwargs
+        @workflow_context.params.kwargs.dup
+      end
+
+      #: () -> Pathname
+      def tmpdir
+        Pathname.new(@workflow_context.tmpdir).realpath
+      end
+    end
+  end
+end