omni_agent 0.1.2 → 0.1.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 8c2b72de1f5d07816381e196c20dec88c00746471223afe27b077aea31952794
-  data.tar.gz: b92ea50f19a39af83fd0d3962b03322e955927fa25c1757836cc5710be8a5be6
+  metadata.gz: '0094c9758c242364af55eeffd4e87143f78b1db9ffabb54c06aaa23234030382'
+  data.tar.gz: 866bc34809e4c1ecbf21679a2d4dab00a0e1383ae3ac3bb8f6f94348d7ae7b6b
 SHA512:
-  metadata.gz: 22d39652f648b750b5aaf292689effe3a855048c7cb4055cdc3f4d21df54611fb16e80d92011a7f893132209013ea7ea222c6da91abe699793529285adc1613a
-  data.tar.gz: d38d01d15b4b767dd6d704bf4f0f8f07d8c69ac52976f3e05d3e5ce66f74f316bfa28f35c30aabd1095b7ba99f8d7c722f6009f85b014fed8884b83df32e292d
+  metadata.gz: bf56557c378dce28538199dc6cf9795923acd7db6ea291dae542845f9f66d923a40ef78e0c1dfdff443e675425c431a58d0cb0d825edc69a171d6b473a0201a1
+  data.tar.gz: a4ee13697bbff34e8a3e0e1f196b282bc05d35ba680b8f19c2860c44527d3b61ed7210ed5a6d756dfcf7acea6bcb4219696d04528b7c3d57a479963f390fc1e4

data/CHANGELOG.md

@@ -22,4 +22,16 @@ All notable changes to this project will be documented in this file.
 ## [0.1.2] - 2026-06-11
 
 ### Fixed
-- Fixed Zeitwerk issue with `OmniAgent::Errors`
+- Fixed Zeitwerk issue with `OmniAgent::Errors`
+
+## [0.1.3] - 2026-06-12
+
+### Added
+- Added `:mock` provider to be able to test `after_generation` and `before_generation` behaviour
+- Added new DSL for tools `stops_generation` which signals to stop the chat flow once this tool is executed
+- Added new instance methods for tools `stops_generation!` and `stops_generation?` that allow to stop the generation within the `execute` logic
+- Added `raw_request` to the agent response, it includes the exact requests params sent to the provider
+- Added `raw_response` to the agent response, it includes the exact response from the provider.
+- Added `generated_messages` to the agent response, it contains all the messages generated by the agent and the user request.
+- Added `history` as a context variable which are the messages to be sent to the provider.
+- Added support for nested datatypes in the `array` schema for tools

data/lib/omni_agent/agent.rb

@@ -147,14 +147,16 @@ module OmniAgent
 
     def run(input, context: {}, prompt_method: nil)
       context = @default_context.merge(context || {})
+      bind_context_instance_variables(context)
 
-      messages = [
-        { role: "system", content: nil },
-        { role: "user", content: input }
-      ]
+      messages = []
 
       run_before_generation_callbacks(input: input, context: context, messages: messages)
+      sync_context_from_instance_variables(context)
+
+      messages.replace(build_messages(input: input, history: context[:history]))
       messages[0][:content] = system_prompt(context: context, prompt_method: prompt_method)
+      initial_messages_count = messages.length - 1
 
       filtered_tools = tool_filter(tools: available_tools, agent_tags: self.class.tags)
 
@@ -163,24 +165,29 @@ module OmniAgent
 
         if response.content && !response.tool_calls?
           messages << { role: "assistant", content: response.content }
+          set_after_generation_state(response: response, messages: messages, initial_messages_count: initial_messages_count)
           run_after_generation_callbacks(input: input, context: context, messages: messages, response: response)
-          return response.content
+          sync_context_from_instance_variables(context)
+          return response
         end
 
-        messages << {
-          role: "assistant",
-          content: response.content,
-          tool_calls: response.raw_response.dig("choices", 0, "message", "tool_calls")
-        }
+        messages << build_assistant_tool_call_message(response)
+        should_stop_generation = false
 
         response.tool_calls.each do |tool_call|
           tool_name = tool_call[:name]
           tool_args = tool_call[:arguments]
           tool_id   = tool_call[:id]
 
-          tool_class = filtered_tools.find { |t| t.name.demodulize == tool_name }
+          tool_class = filtered_tools.find do |t|
+            class_name = t.name.to_s
+            simple_name = class_name.respond_to?(:demodulize) ? class_name.demodulize : class_name.split("::").last
+            simple_name == tool_name
+          end
 
           if tool_class
+            tool_instance = tool_class.new
+
             begin
               result = tool_class.invoke(tool_args)
 
@@ -198,6 +205,9 @@ module OmniAgent
                 content: "Error executing tool: #{e.message}"
               }
             end
+
+            should_stop_generation ||= tool_class.respond_to?(:stops_generation?) && tool_class.stops_generation?
+            should_stop_generation ||= tool_instance.respond_to?(:stops_generation?) && tool_instance.stops_generation?
           else
             messages << {
               role: "tool",
@@ -207,6 +217,13 @@ module OmniAgent
             }
           end
         end
+
+        if should_stop_generation
+          set_after_generation_state(response: response, messages: messages, initial_messages_count: initial_messages_count)
+          run_after_generation_callbacks(input: input, context: context, messages: messages, response: response)
+          sync_context_from_instance_variables(context)
+          return response
+        end
       end
     end
 
@@ -226,6 +243,61 @@ module OmniAgent
       tools
     end
 
+    def build_messages(input:, history:)
+      normalized_history = normalize_history_messages(history)
+
+      [
+        { role: "system", content: nil },
+        *normalized_history,
+        { role: "user", content: input }
+      ]
+    end
+
+    def normalize_history_messages(history)
+      return [] if history.nil?
+      return [] unless history.is_a?(Array)
+
+      history.filter_map do |message|
+        next unless message.is_a?(Hash)
+
+        normalized = message.transform_keys(&:to_sym)
+        normalized.compact
+      end
+    end
+
+    def build_assistant_tool_call_message(response)
+      message = { role: "assistant" }
+      message[:content] = response.content unless response.content.nil?
+
+      raw_tool_calls = response.raw_response.is_a?(Hash) ? response.raw_response.dig("choices", 0, "message", "tool_calls") : nil
+      normalized_tool_calls = normalize_tool_calls_for_message(raw_tool_calls, response.tool_calls)
+      message[:tool_calls] = normalized_tool_calls unless normalized_tool_calls.empty?
+
+      message
+    end
+
+    def normalize_tool_calls_for_message(raw_tool_calls, parsed_tool_calls)
+      return raw_tool_calls if raw_tool_calls.is_a?(Array) && !raw_tool_calls.empty?
+
+      Array(parsed_tool_calls).map do |tool_call|
+        {
+          "id" => tool_call[:id],
+          "type" => "function",
+          "function" => {
+            "name" => tool_call[:name],
+            "arguments" => serialize_tool_arguments(tool_call[:arguments])
+          }
+        }
+      end
+    end
+
+    def serialize_tool_arguments(arguments)
+      return arguments if arguments.is_a?(String)
+      return JSON.generate(arguments) if defined?(JSON)
+
+      arguments.to_s
+    end
+
     def run_alias_entrypoint_logic(alias_name, fallback_input: nil)
       @message = nil
       result = public_send(alias_name)
@@ -249,13 +321,24 @@ module OmniAgent
     end
 
     def run_after_generation_callbacks(input:, context:, messages:, response:)
-      payload = { input: input, context: context, messages: messages, response: response }
+      payload = {
+        input: input,
+        context: context,
+        messages: messages,
+        generated_messages: response.generated_messages,
+        response: response
+      }
 
       self.class.configured_after_generation_callbacks.each do |callback_name|
         invoke_generation_callback(callback_name, payload)
       end
     end
 
+    def set_after_generation_state(response:, messages:, initial_messages_count:)
+      generated_messages = messages.drop(initial_messages_count)
+      @response = response.with_generated_messages(generated_messages)
+    end
+
     def invoke_generation_callback(callback_name, payload)
       original_callback_name = "__omni_agent_original_#{callback_name}".to_sym
       callback_target = if respond_to?(original_callback_name, true)
@@ -265,12 +348,45 @@ module OmniAgent
       end
 
       callback_method = self.class.instance_method(callback_target)
+      bind_context_instance_variables(payload[:context])
 
       if callback_method.arity == 0
         __send__(callback_target)
       else
         __send__(callback_target, payload)
       end
+
+      sync_context_from_instance_variables(payload[:context])
+    end
+
+    def bind_context_instance_variables(context)
+      return unless context.is_a?(Hash)
+
+      @__omni_agent_context_bindings ||= {}
+
+      context.each do |key, value|
+        ivar_name = "@#{key}"
+        next unless ivar_name.match?(/\A@[a-zA-Z_]\w*\z/)
+
+        ivar = ivar_name.to_sym
+        instance_variable_set(ivar, value)
+        @__omni_agent_context_bindings[ivar] = key
+      end
+    end
+
+    def sync_context_from_instance_variables(context)
+      return unless context.is_a?(Hash)
+
+      if instance_variable_defined?(:@__omni_agent_context_bindings)
+        @__omni_agent_context_bindings.each do |ivar, key|
+          next unless instance_variable_defined?(ivar)
+          context[key] = instance_variable_get(ivar)
+        end
+      end
+
+      if instance_variable_defined?(:@history)
+        context[:history] = @history
+      end
     end
 
     def system_prompt(context:, prompt_method: nil)
@@ -292,7 +408,7 @@ module OmniAgent
         isolated_scope.instance_variable_set("@#{key}", value)
       end
 
-      internal_vars = [:@provider, :@chat_options] # Blacklists internal instance variables
+      internal_vars = [ :@provider, :@chat_options ] # Blacklists internal instance variables
 
       (instance_variables - internal_vars).each do |ivar|
         isolated_scope.instance_variable_set(ivar, instance_variable_get(ivar))
@@ -301,7 +417,7 @@ module OmniAgent
       base_prompt = render_prompt_template(base_file_path, isolated_scope)
       method_prompt = render_prompt_template(method_file_path, isolated_scope)
 
-      prompts = [base_prompt, method_prompt].compact.reject(&:empty?)
+      prompts = [ base_prompt, method_prompt ].compact.reject(&:empty?)
       return prompts.join("\n\n") if prompts.any?
 
       "You are a helpful assistant with access to local tools."
@@ -325,4 +441,4 @@ module OmniAgent
         .downcase
     end
   end
-end
+end

data/lib/omni_agent/providers/base.rb

@@ -13,6 +13,32 @@ module OmniAgent
         raise NotImplementedError, "Providers must implement #chat"
       end
 
+      def validate_messages!(messages, allowed_roles:)
+        unless messages.is_a?(Array)
+          raise OmniAgent::Error, "messages must be an array"
+        end
+
+        normalized_roles = Array(allowed_roles).map(&:to_s)
+
+        messages.each_with_index do |message, index|
+          unless message.is_a?(Hash)
+            raise OmniAgent::Error, "message at index #{index} must be a hash"
+          end
+
+          role = message[:role] || message["role"]
+          if role.nil?
+            raise OmniAgent::Error, "message at index #{index} is missing role"
+          end
+
+          role_name = role.to_s
+          unless normalized_roles.include?(role_name)
+            raise OmniAgent::Error, "invalid message role '#{role_name}' at index #{index}"
+          end
+
+          validate_message_payload!(message, index, role_name)
+        end
+      end
+
       protected
 
       def default_api_key
@@ -22,6 +48,33 @@ module OmniAgent
       def default_model
         raise NotImplementedError, "Providers must define a default model"
       end
+
+      def validate_message_payload!(message, index, role_name)
+        has_content = message.key?(:content) || message.key?("content")
+        content_value = message[:content] || message["content"]
+
+        case role_name
+        when "assistant"
+          has_tool_calls = message.key?(:tool_calls) || message.key?("tool_calls")
+          return if has_content || has_tool_calls
+
+          raise OmniAgent::Error, "assistant message at index #{index} must include content or tool_calls"
+        when "tool"
+          tool_call_id = message[:tool_call_id] || message["tool_call_id"]
+          tool_name = message[:name] || message["name"]
+          unless tool_call_id && tool_name && has_content
+            raise OmniAgent::Error, "tool message at index #{index} must include tool_call_id, name, and content"
+          end
+        else
+          unless has_content
+            raise OmniAgent::Error, "#{role_name} message at index #{index} must include content"
+          end
+        end
+
+        if has_content && content_value.nil?
+          raise OmniAgent::Error, "message content at index #{index} cannot be nil"
+        end
+      end
     end
   end
-end
+end

data/lib/omni_agent/providers/mock.rb

@@ -0,0 +1,40 @@
+module OmniAgent
+  module Providers
+    class Mock < Base
+      LOREM_IPSUM = "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
+
+      def chat(messages:, tools: [], **_options)
+        validate_messages!(messages, allowed_roles: %i[system user assistant tool])
+
+        OmniAgent::Providers::Response.new(
+          content: LOREM_IPSUM,
+          raw_request: {
+            model: model,
+            messages: messages,
+            tools: tools
+          },
+          raw_response: {
+            "choices" => [
+              {
+                "message" => {
+                  "content" => LOREM_IPSUM
+                }
+              }
+            ]
+          },
+          tool_calls: []
+        )
+      end
+
+      protected
+
+      def default_api_key
+        nil
+      end
+
+      def default_model
+        "mock"
+      end
+    end
+  end
+end

data/lib/omni_agent/providers/openai.rb

@@ -1,18 +1,40 @@
 # lib/omni_agent/providers/openai.rb
-require 'json'
+require "json"
 
 module OmniAgent
   module Providers
     class OpenAI < Base
       begin
-        require 'openai'
+        require "openai"
       rescue LoadError
-        raise OmniAgent::MissingDependencyError, 
+        raise OmniAgent::MissingDependencyError,
               "The 'openai' gem is required to use the OpenAI provider. " \
               "Please add `gem 'openai'` to your Gemfile."
       end
-        
+
       def chat(messages:, tools: [], **options)
+        validate_messages!(messages, allowed_roles: %i[system user assistant tool])
+
+        messages = messages.map do |msg|
+          clean_msg = msg.reject { |_, v| v.nil? }
+
+          if !clean_msg.key?(:tool_calls) && !clean_msg[:content].is_a?(String)
+            clean_msg[:content] = ""
+          end
+
+          if clean_msg[:tool_calls].is_a?(Array)
+            clean_msg[:tool_calls].each do |tc|
+              if tc.dig("function", "arguments").is_a?(Hash) || tc.dig(:function, :arguments).is_a?(Hash)
+                func = tc["function"] || tc[:function]
+                args_key = func.key?("arguments") ? "arguments" : :arguments
+                func[args_key] = func[args_key].to_json
+              end
+            end
+          end
+
+          clean_msg
+        end
+
         openai_tools = tools.map { |tool| format_tool(tool) }
 
         payload = {
@@ -24,9 +46,9 @@ module OmniAgent
 
         response = client.chat.completions.create(**payload)
 
-        parse_response(response)
-      rescue => e 
-        puts "Error during OpenAI chat: #{e.message}"
+        parse_response(response, payload)
+      rescue => e
+        raise OmniAgent::Error, "Error during OpenAI chat: #{e.message}"
       end
 
       protected
@@ -36,7 +58,7 @@ module OmniAgent
       end
 
       def default_api_key
-        ENV.fetch('OPENAI_ACCESS_TOKEN', nil)
+        ENV.fetch("OPENAI_ACCESS_TOKEN", nil)
       end
 
       def default_model
@@ -46,20 +68,25 @@ module OmniAgent
       private
 
       def format_tool(tool_class)
+        schema = tool_class.json_schema.dup
+
+        schema[:additionalProperties] = false
+
         {
           type: "function",
           function: {
             name: tool_class.name.split("::").last,
             description: tool_class.description,
-            parameters: tool_class.json_schema 
+            parameters: schema
           }
         }
       end
 
-      def parse_response(raw_response)
-        choices = raw_response.respond_to?(:choices) ? raw_response.choices : raw_response["choices"]
+      def parse_response(raw_response, raw_request)
+        provider_raw_response = raw_response.respond_to?(:to_h) ? raw_response.to_h : raw_response
+        choices = raw_response.respond_to?(:choices) ? raw_response.choices : provider_raw_response["choices"]
         first_choice = choices&.first || {}
-        
+
         message = first_choice.respond_to?(:message) ? first_choice.message : (first_choice["message"] || {})
         content = message.respond_to?(:content) ? message.content : message["content"]
 
@@ -75,35 +102,13 @@ module OmniAgent
           }
         end
 
-        compat_tool_calls = raw_tool_calls.map do |tc|
-          fn = tc.respond_to?(:function) ? tc.function : tc["function"]
-          {
-            "id" => tc.respond_to?(:id) ? tc.id : tc["id"],
-            "type" => "function",
-            "function" => {
-              "name" => fn.respond_to?(:name) ? fn.name : fn["name"],
-              "arguments" => fn.respond_to?(:arguments) ? fn.arguments : fn["arguments"]
-            }
-          }
-        end.compact
-
-        compat_raw_response = {
-          "choices" => [
-            {
-              "message" => {
-                "content" => content,
-                "tool_calls" => (compat_tool_calls.empty? ? nil : compat_tool_calls)
-              }.compact
-            }
-          ]
-        }
-
         OmniAgent::Providers::Response.new(
           content: content,
-          raw_response: compat_raw_response,
+          raw_response: provider_raw_response,
+          raw_request: raw_request,
           tool_calls: tool_calls
         )
       end
     end
   end
-end
+end

data/lib/omni_agent/providers/response.rb

@@ -1,17 +1,28 @@
 module OmniAgent
   module Providers
     class Response
-      attr_reader :content, :tool_calls, :raw_response
+      attr_reader :content, :tool_calls, :raw_response, :raw_request, :generated_messages
 
-      def initialize(content:, tool_calls: [], raw_response: nil)
+      def initialize(content:, tool_calls: [], raw_response: nil, raw_request: nil, generated_messages: [])
         @content = content
         @tool_calls = tool_calls || []
         @raw_response = raw_response
+        @raw_request = raw_request
+        @generated_messages = generated_messages || []
       end
 
       def tool_calls?
         @tool_calls.any?
       end
+
+      def answer
+        @content
+      end
+
+      def with_generated_messages(messages)
+        @generated_messages = Array(messages)
+        self
+      end
     end
   end
-end
+end

data/lib/omni_agent/providers.rb

@@ -3,8 +3,9 @@ module OmniAgent
   module Providers
     def self.registry
       {
-        openai: OmniAgent::Providers::OpenAI 
+        openai: OmniAgent::Providers::OpenAI,
+        mock: OmniAgent::Providers::Mock
       }
     end
   end
-end
+end

data/lib/omni_agent/tool/schema_builder.rb

@@ -20,10 +20,24 @@ module OmniAgent
         add_property(name, type: "boolean", description: description, required: required)
       end
 
-      def array(name, items_type:, description: nil, required: true)
-        property = { type: "array", items: { type: items_type } }
+      def array(name, items_type: nil, description: nil, required: true, &block)
+        property = { type: "array" }
         property[:description] = description if description
-        
+
+        if block_given?
+          nested_builder = SchemaBuilder.new
+          nested_builder.instance_eval(&block)
+
+          property[:items] = {
+            type: "object",
+            properties: nested_builder.properties,
+            required: nested_builder.required_fields,
+            additionalProperties: false
+          }
+        else
+          property[:items] = { type: items_type || "string" }
+        end
+
         @properties[name] = property
         @required_fields << name.to_s if required
       end
@@ -35,7 +49,7 @@ module OmniAgent
         if block_given?
           nested_builder = SchemaBuilder.new
           nested_builder.instance_eval(&block)
-          
+
           property[:properties] = nested_builder.properties
           property[:required] = nested_builder.required_fields
           property[:additionalProperties] = false
@@ -52,10 +66,10 @@ module OmniAgent
       def add_property(name, type:, description:, required:)
         property = { type: type }
         property[:description] = description if description
-        
+
         @properties[name] = property
         @required_fields << name.to_s if required
       end
     end
   end
-end
+end

data/lib/omni_agent/tool.rb

@@ -1,4 +1,3 @@
-# lib/omni_agents/tool.rb
 module OmniAgent
   class Tool
     class << self
@@ -26,7 +25,7 @@ module OmniAgent
         if block_given?
           builder = SchemaBuilder.new
           builder.instance_eval(&block)
-          
+
           @properties = builder.properties
           @required = builder.required_fields
         end
@@ -37,13 +36,26 @@ module OmniAgent
           type: "object",
           properties: @properties || {},
           required: @required || [],
-          additionalProperties: false 
+          additionalProperties: false
         }
       end
 
       def invoke(arguments_hash)
         kwargs = arguments_hash.transform_keys(&:to_sym)
-        new.execute(**kwargs)
+
+        valid_keys = (@properties || {}).keys.map(&:to_sym)
+
+        filtered_kwargs = kwargs.slice(*valid_keys)
+
+        new.execute(**filtered_kwargs)
+      end
+
+      def stops_generation(value = true)
+        @stops_generation = !!value
+      end
+
+      def stops_generation?
+        @stops_generation == true
       end
 
       private
@@ -61,8 +73,16 @@ module OmniAgent
       end
     end
 
+    def stop_generation!
+      @stops_generation_instance = true
+    end
+
+    def stops_generation?
+      @stops_generation_instance == true
+    end
+
     def execute(**args)
       raise NotImplementedError, "#{self.class.name} must implement #execute"
     end
   end
-end
+end

data/lib/omni_agent/version.rb

@@ -1,3 +1,3 @@
 module OmniAgent
-  VERSION = "0.1.2"
+  VERSION = "0.1.3"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: omni_agent
 version: !ruby/object:Gem::Version
-  version: 0.1.2
+  version: 0.1.3
 platform: ruby
 authors:
 - ACR1209
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2026-06-11 00:00:00.000000000 Z
+date: 2026-06-13 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rails
@@ -67,6 +67,7 @@ files:
 - lib/omni_agent/errors.rb
 - lib/omni_agent/providers.rb
 - lib/omni_agent/providers/base.rb
+- lib/omni_agent/providers/mock.rb
 - lib/omni_agent/providers/openai.rb
 - lib/omni_agent/providers/response.rb
 - lib/omni_agent/tasks/omni_agent_tasks.rake