ruboty-ai_agent 0.1.0

data/lib/ruboty/ai_agent/mcp_configuration.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Ruboty
+  module AiAgent
+    # @rbs!
+    #   type transports = :http | :websocket
+
+    McpConfiguration = Data.define(
+      :name, #: String
+      :transport, #: transports
+      :headers, #: Hash[String, String]
+      :url #: String
+    )
+
+    # Save MCP configuration details.
+    class McpConfiguration
+      include Recordable
+
+      # @rbs name: String
+      # @rbs transport: transports
+      # @rbs url: String
+      # @rbs headers: Hash[String, String]?
+      def initialize(name:, transport:, url:, headers: {})
+        # No superclass method `initialize` in RBS.
+        super(name:, transport:, headers:, url:) # steep:ignore UnexpectedKeywordArgument
+      end
+
+      register_record_type :mcp_configuration
+    end
+  end
+end

data/lib/ruboty/ai_agent/record_set.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module Ruboty
+  module AiAgent
+    # A set of records.
+    # @rbs generic Record
+    class RecordSet
+      attr_reader :database #: Ruboty::AiAgent::Database
+
+      def initialize(database:)
+        @database = database
+      end
+
+      def namespace_keys #: Array[Database::keynable]
+        raise NotImplementedError, 'Subclasses must implement the namespace_keys method'
+      end
+
+      def length #: Integer
+        database.len(*namespace_keys)
+      end
+
+      def all #: untyped
+        database.fetch(*namespace_keys)
+      end
+
+      def all_values #: Array[Record]
+        case (kv = all)
+        when Hash
+          kv.values
+        when Array
+          kv
+        else
+          []
+        end
+      end
+
+      def keys #: Array[Database::keynable]
+        database.keys(*namespace_keys)
+      end
+
+      # @rbs key: Database::keynable
+      # @rbs return: Record | nil
+      def fetch(key)
+        database.fetch(*namespace_keys, key)
+      end
+
+      # @rbs key: Database::keynable
+      # @rbs record: Record
+      # @rbs return: void
+      def store(record, key:)
+        database.store(record, at: [*namespace_keys, key])
+      end
+
+      # @rbs key: Database::keynable
+      # @rbs return: void
+      def remove(key)
+        database.delete(*namespace_keys, key)
+      end
+
+      # @rbs key: Database::keynable
+      # @rbs return: boolish
+      def key?(key)
+        database.key?(*namespace_keys, key)
+      end
+
+      def clear #: void
+        database.delete(*namespace_keys)
+      end
+    end
+  end
+end

data/lib/ruboty/ai_agent/recordable.rb

@@ -0,0 +1,116 @@
+# frozen_string_literal: true
+
+module Ruboty
+  module AiAgent
+    # @rbs!
+    #   interface _WithToH
+    #     def to_h: () -> Hash[Database::keynable, untyped]
+    #   end
+
+    # Convertable between Hash and Recordable bidirectionally.
+    # @rbs module-self _WithToH
+    module Recordable
+      class << self
+        def included(base)
+          base.extend(ClassMethods)
+          base.prepend(PrependMethods)
+        end
+
+        # @rbs @record_types: Hash[Symbol, Class]
+
+        def record_types #: Hash[Symbol, Class]
+          @record_types ||= {}
+        end
+
+        # @rbs hash: Hash[Symbol, untyped]?
+        # @rbs return: bool
+        def convertable?(hash)
+          return false unless hash.is_a?(Hash)
+
+          type = hash[:record_type]
+          type && record_types.include?(type)
+        end
+
+        # @rbs value: untyped
+        # @rbs return: untyped
+        def instantiate_recursively(value)
+          case value
+          when Hash
+            transformed = value.transform_values { |v| instantiate_recursively(v) }
+            if convertable?(transformed)
+              record_from_hash(transformed)
+            else
+              transformed
+            end
+          when Array
+            value.map { |v| instantiate_recursively(v) }
+          else
+            value
+          end
+        end
+
+        # @rbs value: untyped
+        # @rbs return: untyped
+        def hashify_recursively(value)
+          case value
+          when Recordable
+            hashify_recursively(value.to_h)
+          when Hash
+            value.transform_values { |v| hashify_recursively(v) }
+          when Array
+            value.map { |v| hashify_recursively(v) }
+          else
+            value
+          end
+        end
+
+        # @rbs record: Recordable
+        # @rbs return: Hash[Database::keynable, untyped]
+        def record_to_hash(record)
+          record.to_h
+        end
+
+        # @rbs hash: Hash[Symbol, untyped]
+        # @rbs return: Recordable
+        def record_from_hash(hash)
+          type = hash[:record_type]
+          klass = record_types[type]
+          raise "Unknown record type: #{type}" unless klass
+
+          klass.new(**hash.except(:record_type))
+        end
+      end
+
+      # @rbs module-self Class
+      module ClassMethods
+        attr_accessor :record_type #: Symbol
+
+        # @rbs name: Symbol
+        def register_record_type(name)
+          name = name.to_sym
+          self.record_type = name
+
+          Recordable.record_types.merge!({ name => self }) do
+            raise "Duplicate record type: #{name}"
+          end
+        end
+      end
+
+      # @rbs module-self Recordable::ClassMethods.instance
+      module PrependMethods
+        def to_h #: Hash[Database::keynable, untyped]
+          {
+            record_type: record_type,
+            **super
+          }
+        end
+      end
+
+      # @rbs %a{pure}
+      def record_type #: Symbol
+        self.class #: singleton(::Object) & ClassMethods
+            .record_type
+      end
+    end
+  end
+end

data/lib/ruboty/ai_agent/token_usage.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module Ruboty
+  module AiAgent
+    # Token usage information from LLM responses
+    class TokenUsage
+      include Recordable
+
+      register_record_type :token_usage
+
+      attr_reader :prompt_tokens #: Integer
+      attr_reader :completion_tokens #: Integer
+      attr_reader :total_tokens #: Integer
+      attr_reader :token_limit #: Integer?
+
+      # @rbs prompt_tokens: Integer
+      # @rbs completion_tokens: Integer
+      # @rbs total_tokens: Integer
+      # @rbs ?token_limit: Integer?
+      def initialize(prompt_tokens:, completion_tokens:, total_tokens:, token_limit: nil)
+        @prompt_tokens = prompt_tokens
+        @completion_tokens = completion_tokens
+        @total_tokens = total_tokens
+        @token_limit = token_limit
+      end
+
+      # Calculate usage percentage if token limit is available
+      # @rbs return: (Float | nil)
+      def usage_percentage
+        return nil unless token_limit
+
+        (total_tokens.to_f / token_limit * 100).round(2).to_f
+      end
+
+      def to_h #: Hash[Symbol, untyped]
+        {
+          prompt_tokens:,
+          completion_tokens:,
+          total_tokens:,
+          token_limit:
+        }
+      end
+    end
+  end
+end

data/lib/ruboty/ai_agent/tool.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Ruboty
+  module AiAgent
+    # Define a tool that the AI agent can use.
+    class Tool
+      attr_reader :name, :title, :description #: String
+      attr_reader :input_schema #: Hash[untyped, untyped]?
+      attr_reader :on_call #: (^(Hash[String, untyped]) -> String )?
+
+      # @rbs name: String
+      # @rbs title: String
+      # @rbs description: String
+      # @rbs input_schema: Hash[untyped, untyped]?
+      # @rbs &on_call: ? (Hash[String, untyped]) -> String
+      def initialize(name:, title:, description:, input_schema:, &on_call) #: void
+        @name = name
+        @title = title
+        @description = description
+        @input_schema = input_schema
+        @on_call = on_call
+      end
+
+      def call(params) #: String?
+        on_call&.call(params)
+      end
+    end
+  end
+end

data/lib/ruboty/ai_agent/user.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module Ruboty
+  module AiAgent
+    # User class to manage user-specific data.
+    class User
+      attr_reader :database #: Ruboty::AiAgent::Database
+      attr_reader :id #: String
+
+      class << self
+        def find_or_create(database:, id:) #: User
+          new(database: database, id: id)
+        end
+      end
+
+      def initialize(database:, id:)
+        @database = database
+        @id = id
+      end
+
+      # @rbs %a{memorized}
+      def mcp_configurations #: UserMcpConfigurations
+        @mcp_configurations ||= UserMcpConfigurations.new(database: database, user_id: id)
+      end
+
+      def mcp_clients #: Array[UserMcpClient]
+        mcp_configurations.all_values.map do |config|
+          UserMcpClient.new(user: self, mcp_name: config.name)
+        end
+      end
+
+      # @rbs %a{memorized}
+      def ai_memories #: UserAiMemories
+        @ai_memories ||= UserAiMemories.new(database: database, user_id: id)
+      end
+
+      def system_prompt #: String?
+        database.fetch(:users, id, :system_prompt)
+      end
+
+      # @rbs prompt: String?
+      def system_prompt=(prompt)
+        database.store(prompt, at: [:users, id, :system_prompt])
+      end
+
+      # @rbs %a{memorized}
+      def mcp_tools_caches #: UserMcpToolsCaches
+        @mcp_tools_caches ||= UserMcpToolsCaches.new(database: database, user_id: id)
+      end
+    end
+  end
+end

data/lib/ruboty/ai_agent/user_ai_memories.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Ruboty
+  module AiAgent
+    # Manage Ai Memories for a specific user.
+    class UserAiMemories < UserAssociations #[String]
+      self.association_key = :ai_memories
+
+      # @rbs memory: String
+      def add(memory) #: void
+        next_id = ((keys.map { |key| Integer(key) }.max || 0) + 1).to_s # steep:ignore
+        store(memory, key: next_id)
+        next_id
+      end
+    end
+  end
+end

data/lib/ruboty/ai_agent/user_associations.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module Ruboty
+  module AiAgent
+    # A set of records for a specific user.
+    # @rbs generic Record
+    class UserAssociations < RecordSet #[Record]
+      attr_reader :user_id #: String
+
+      def initialize(database:, user_id:)
+        super(database:)
+
+        @user_id = user_id
+      end
+
+      # @rbs!
+      #   def self.association_key: () -> Symbol
+      #   def self.association_key=: (Symbol) -> Symbol
+
+      # @rbs skip
+      class << self
+        attr_accessor :association_key
+      end
+
+      def association_key #: Symbol
+        self.class.association_key || raise(NotImplementedError, 'Subclasses must set the association_key method')
+      end
+
+      def namespace_keys
+        [:users, user_id, association_key]
+      end
+    end
+  end
+end

data/lib/ruboty/ai_agent/user_mcp_caches.rb

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+module Ruboty
+  module AiAgent
+    # Manage MCP caches for a specific user.
+    # @abstract
+    # @rbs generic D < Object
+    class UserMcpCaches < UserAssociations #[CachedValue[D]]
+      # @rbs!
+      #   def self.cache_type_key: () -> Symbol?
+      #   def self.cache_type_key=: (Symbol?) -> Symbol?
+
+      # @rbs!
+      #   def self.cache_duration: () -> Integer?
+      #   def self.cache_duration=: (Integer?) -> Integer?
+
+      class << self
+        # @rbs skip
+        attr_accessor :cache_type_key
+        # @rbs skip
+        attr_accessor :cache_duration
+
+        def association_key #: Symbol
+          raise 'Subclasses must set the cache_type_key method' unless cache_type_key
+
+          :"mcp_caches:#{cache_type_key}"
+        end
+      end
+
+      def cache_type_key #: Symbol
+        self.class.cache_type_key || raise(NotImplementedError, 'Subclasses must set the cache_type_key method')
+      end
+
+      def cache_duration #: Integer
+        self.class.cache_duration || raise(NotImplementedError, 'Subclasses must set the cache_duration method')
+      end
+
+      # @rbs key: String)
+      # @rbs &block: () -> (D | CachedValue[D])
+      # @rbs return: D
+      def fetch_or_store_data(key, &block)
+        if (fetched = fetch_data(key))
+          fetched
+        else
+          data = block.call
+          cache = data.is_a?(CachedValue) ? data : CachedValue.new(data:, expires_at: Time.now + cache_duration)
+          store(cache, key:)
+          cache.data
+        end
+      end
+
+      # @rbs key: String
+      # @rbs return: D?
+      def fetch_data(key)
+        fetch(key)&.data
+      end
+
+      # @rbs override
+      def all
+        super().reject { |(_key, cache)| cache.expired? }
+      end
+
+      # @rbs override
+      def fetch(key)
+        cache = super(key)
+
+        if cache&.expired?
+          remove(key)
+          nil
+        end
+
+        cache
+      end
+
+      # @rbs key: String
+      # @rbs data: D
+      # @rbs return: void
+      def store_data(data, key:)
+        expires_at = Time.now + cache_duration
+
+        cache = CachedValue.new(
+          data:,
+          expires_at:
+        ) #: CachedValue[D]
+
+        store(cache, key:)
+      end
+    end
+  end
+end

data/lib/ruboty/ai_agent/user_mcp_client.rb

@@ -0,0 +1,93 @@
+# frozen_string_literal: true
+
+module Ruboty
+  module AiAgent
+    # Wrapper for MCP client with caching support for a specific user
+    class UserMcpClient
+      attr_reader :user #: User
+      attr_reader :mcp_name #: String
+
+      # @rbs user: User
+      # @rbs mcp_name: String
+      # @rbs return: void
+      def initialize(user:, mcp_name:)
+        @user = user
+        @mcp_name = mcp_name
+      end
+
+      # @rbs return: Array[tool_def]
+      def list_tools
+        user.mcp_tools_caches.fetch_or_store_data(mcp_name) do
+          mcp_client.list_tools
+        end
+      end
+
+      # @rbs name: String
+      # @rbs arguments: Hash[String, untyped]
+      # @rbs &block: ? (Hash[String, untyped]) -> void
+      # @rbs return: untyped
+      def call_tool(name, arguments = {}, &block)
+        mcp_client.call_tool(name, arguments, &block)
+      end
+
+      # @rbs return: untyped
+      def list_prompts
+        mcp_client.list_prompts
+      end
+
+      # @rbs name: String
+      # @rbs arguments: Hash[String, untyped]
+      # @rbs return: untyped
+      def get_prompt(name, arguments = {})
+        mcp_client.get_prompt(name, arguments)
+      end
+
+      # @rbs return: untyped
+      def list_resources
+        mcp_client.list_resources
+      end
+
+      # @rbs uri: String
+      # @rbs return: untyped
+      def read_resource(uri)
+        mcp_client.read_resource(uri)
+      end
+
+      # @rbs return: untyped
+      def ping
+        mcp_client.ping
+      end
+
+      # @rbs return: untyped
+      def initialize_session
+        mcp_client.initialize_session
+      end
+
+      # @rbs return: untyped
+      def cleanup_session
+        mcp_client.cleanup_session
+      end
+
+      private
+
+      # @rbs return: McpConfiguration
+      def configuration
+        user.mcp_configurations.all_values.find { |config| config.name == mcp_name } ||
+          raise("MCP configuration not found: #{mcp_name}")
+      end
+
+      # @rbs return: HttpMcpClient
+      def mcp_client
+        @mcp_client ||= case configuration.transport
+                        when :http
+                          HttpMcpClient.new(
+                            url: configuration.url,
+                            headers: configuration.headers || {}
+                          )
+                        else
+                          raise "Unknown MCP server type: #{configuration.transport}"
+                        end
+      end
+    end
+  end
+end

data/lib/ruboty/ai_agent/user_mcp_configurations.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Ruboty
+  module AiAgent
+    # Manage MCP configurations for a specific user.
+    class UserMcpConfigurations < UserAssociations #[McpConfiguration]
+      self.association_key = :mcp_configurations
+
+      # @rbs mcp_configuration: McpConfiguration:
+      def add(mcp_configuration) #: void
+        store(mcp_configuration, key: mcp_configuration.name)
+      end
+    end
+  end
+end

data/lib/ruboty/ai_agent/user_mcp_tools_caches.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module Ruboty
+  module AiAgent
+    # @rbs!
+    #   type tool_def = Hash[Symbol | String, untyped]
+
+    # Manage tools cache for a specific user and server.
+    class UserMcpToolsCaches < UserMcpCaches #[Array[tool_def]]
+      self.cache_type_key = :tools
+      self.cache_duration = ENV.fetch('MCP_CACHE_DURATION', 600).to_i
+    end
+  end
+end

data/lib/ruboty/ai_agent/version.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module Ruboty
+  module AiAgent
+    VERSION = '0.1.0'
+  end
+end

data/lib/ruboty/ai_agent.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+require 'json'
+
+require 'ruboty/ai_agent/version'
+require 'ruboty/handlers/ai_agent'
+
+module Ruboty
+  # Add AI Agent feature to Ruboty.
+  module AiAgent
+    class Error < StandardError; end
+
+    autoload :Actions, 'ruboty/ai_agent/actions'
+    autoload :Agent, 'ruboty/ai_agent/agent'
+    autoload :ChatMessage, 'ruboty/ai_agent/chat_message'
+    autoload :ChatThread, 'ruboty/ai_agent/chat_thread'
+    autoload :ChatThreadAssociations, 'ruboty/ai_agent/chat_thread_associations'
+    autoload :ChatThreadMessages, 'ruboty/ai_agent/chat_thread_messages'
+    autoload :Commands, 'ruboty/ai_agent/commands'
+    autoload :Database, 'ruboty/ai_agent/database'
+    autoload :GlobalSettings, 'ruboty/ai_agent/global_settings'
+    autoload :HttpMcpClient, 'ruboty/ai_agent/http_mcp_client'
+    autoload :LLM, 'ruboty/ai_agent/llm'
+    autoload :CachedValue, 'ruboty/ai_agent/cached_value'
+    autoload :McpClient, 'ruboty/ai_agent/mcp_client'
+    autoload :McpClients, 'ruboty/ai_agent/mcp_clients'
+    autoload :McpConfiguration, 'ruboty/ai_agent/mcp_configuration'
+    autoload :Recordable, 'ruboty/ai_agent/recordable'
+    autoload :RecordSet, 'ruboty/ai_agent/record_set'
+    autoload :TokenUsage, 'ruboty/ai_agent/token_usage'
+    autoload :Tool, 'ruboty/ai_agent/tool'
+    autoload :User, 'ruboty/ai_agent/user'
+    autoload :UserAiMemories, 'ruboty/ai_agent/user_ai_memories'
+    autoload :UserAssociations, 'ruboty/ai_agent/user_associations'
+    autoload :UserMcpCaches, 'ruboty/ai_agent/user_mcp_caches'
+    autoload :UserMcpClient, 'ruboty/ai_agent/user_mcp_client'
+    autoload :UserMcpConfigurations, 'ruboty/ai_agent/user_mcp_configurations'
+    autoload :UserMcpToolsCaches, 'ruboty/ai_agent/user_mcp_tools_caches'
+  end
+end