rager 0.6.0 → 0.8.0

data/lib/rager/result.rb

@@ -26,7 +26,7 @@ module Rager
     sig { returns(T.nilable(Rager::Types::Output)) }
     attr_reader :output
 
-    sig { returns(Rager::Options) }
+    sig { returns(Rager::Types::Options) }
     attr_reader :options
 
     sig { returns(Integer) }
@@ -35,23 +35,29 @@ module Rager
     sig { returns(Integer) }
     attr_reader :end_time
 
-    sig { returns(T.nilable(String)) }
-    attr_reader :name
+    sig { returns(Integer) }
+    attr_reader :attempt
+
+    sig { returns(T::Array[String]) }
+    attr_reader :errors
+
+    sig { returns(T::Array[String]) }
+    attr_reader :input_ids
 
     sig { returns(T.nilable(String)) }
-    attr_reader :context_name
+    attr_reader :name
 
     sig { returns(T::Array[String]) }
     attr_reader :tags
 
-    sig { returns(T::Array[String]) }
-    attr_reader :input_ids
+    sig { returns(T.nilable(String)) }
+    attr_reader :context_name
 
     sig { returns(T::Array[String]) }
-    attr_reader :errors
+    attr_reader :context_tags
 
-    sig { returns(Integer) }
-    attr_reader :attempt
+    sig { returns(T.untyped) }
+    attr_reader :transform
 
     sig do
       params(
@@ -60,33 +66,20 @@ module Rager
         operation: Rager::Operation,
         input: Rager::Types::Input,
         output: T.nilable(Rager::Types::Output),
-        options: Rager::Options,
+        options: Rager::Types::Options,
         start_time: Integer,
         end_time: Integer,
+        attempt: Integer,
+        errors: T::Array[String],
+        input_ids: T::Array[String],
         name: T.nilable(String),
-        context_name: T.nilable(String),
         tags: T::Array[String],
-        input_ids: T::Array[String],
-        errors: T::Array[String],
-        attempt: Integer
+        context_name: T.nilable(String),
+        context_tags: T::Array[String],
+        transform: T.untyped
       ).void
     end
-    def initialize(
-      id:,
-      context_id:,
-      operation:,
-      input:,
-      output:,
-      options:,
-      start_time:,
-      end_time:,
-      name: nil,
-      context_name: nil,
-      tags: [],
-      input_ids: [],
-      errors: [],
-      attempt: 0
-    )
+    def initialize(id:, context_id:, operation:, input:, output:, options:, start_time:, end_time:, attempt: 0, errors: [], input_ids: [], name: nil, tags: [], context_name: nil, context_tags: [], transform: nil)
       @id = id
       @context_id = context_id
       @operation = operation
@@ -95,12 +88,14 @@ module Rager
       @options = options
       @start_time = start_time
       @end_time = end_time
-      @name = T.let(name, T.nilable(String))
-      @context_name = T.let(context_name, T.nilable(String))
-      @tags = T.let(tags, T::Array[String])
-      @input_ids = T.let(input_ids, T::Array[String])
-      @errors = T.let(errors, T::Array[String])
       @attempt = attempt
+      @errors = T.let(errors, T::Array[String])
+      @input_ids = T.let(input_ids, T::Array[String])
+      @name = name
+      @tags = tags
+      @context_name = context_name
+      @context_tags = T.let(context_tags, T::Array[String])
+      @transform = transform
 
       @stream = T.let(nil, T.nilable(Rager::Types::Stream))
       @buffer = T.let([], Rager::Types::Buffer)
@@ -118,7 +113,7 @@ module Rager
     end
 
     sig { returns(Rager::Types::Output) }
-    def out
+    def stream
       return T.must(@output) unless stream?
       return @buffer.each if @consumed
 
@@ -138,123 +133,98 @@ module Rager
       @stream
     end
 
-    sig { returns(Rager::Types::NonStreamOutput) }
-    def mat
-      return T.cast(@output, Rager::Types::NonStreamOutput) unless stream?
+    sig { returns(T.untyped) }
+    def value
+      return @transform if @transform
 
-      if !@consumed
-        T.cast(out, Rager::Types::Stream).each { |_| }
-      end
+      return T.cast(@output, Rager::Types::NonStreamOutput) unless stream?
 
-      parts = {}
+      T.cast(stream, Rager::Types::Stream).each { |_| } unless @consumed
 
-      @buffer.each do |message_delta|
-        parts[message_delta.index] =
-          (parts[message_delta.index] || "") + message_delta.content
-      end
-
-      parts
-        .sort_by { |index, _| index }
-        .map { |_, content| content }
+      @buffer
+        .group_by(&:index)
+        .transform_values { |deltas| deltas.map(&:content).join }
+        .sort
+        .map(&:last)
         .then { |parts| (parts.length == 1) ? parts.first : parts }
     end
 
     sig do
-      returns(
-        T.any(
-          String,
-          T::Hash[String, T.untyped],
-          T::Array[T.untyped]
-        )
-      )
-    end
-    def serialize_input
-      case @input
-      when String
-        @input
-      when Hash
-        @input.transform_keys(&:to_s)
-      when Array
-        @input.map do |item|
-          if !item.is_a?(String)
-            item.serialize
-          else
-            item
-          end
-        end
-      else
-        @input.to_s
-      end
+      params(
+        block: T.proc.params(
+          value: Rager::Types::NonStreamOutput
+        ).returns(T.untyped)
+      ).returns(Rager::Result)
     end
+    def let(&block)
+      transform = yield(value)
 
-    sig { returns(T.nilable(Rager::Types::NonStreamOutput)) }
-    def serialize_output
-      return nil unless success?
-      return @consumed ? mat : "[STREAM]" if stream?
-      T.cast(@output, Rager::Types::NonStreamOutput)
+      Result.new(
+        id: @id,
+        context_id: @context_id,
+        operation: @operation,
+        input: @input,
+        output: @output,
+        options: @options,
+        start_time: @start_time,
+        end_time: @end_time,
+        attempt: @attempt,
+        errors: @errors,
+        input_ids: @input_ids,
+        name: @name,
+        tags: @tags,
+        context_name: @context_name,
+        context_tags: @context_tags,
+        transform: transform
+      )
     end
 
-    sig { returns(T::Hash[String, T.untyped]) }
+    sig { returns(T::Hash[Symbol, T.untyped]) }
     def to_h
       {
-        id: @id,
         version: "1",
         data: {
+          id: @id,
           context_id: @context_id,
           operation: @operation.serialize,
-          input: {input: serialize_input},
-          output: {output: serialize_output},
+          input: @input,
+          output: display_output,
           options: @options.serialize_safe,
           start_time: @start_time,
           end_time: @end_time,
+          attempt: @attempt,
+          errors: @errors,
+          input_ids: @input_ids,
           name: @name,
-          context_name: @context_name,
           tags: @tags,
-          input_ids: @input_ids,
-          errors: @errors,
-          attempt: @attempt
+          context_name: @context_name,
+          context_tags: @context_tags
         }
       }
     end
 
+    sig { returns(String) }
+    def to_json
+      to_h.to_json
+    end
+
     sig { void }
     def log
-      return unless Rager.config.logger_type
-
-      json = to_h.to_json
-
-      logger = Rager.config.logger
-
-      case Rager.config.logger_type
-      when Rager::Logger::Stdout
-        success? ? logger.info(json) : logger.error(json)
-      when Rager::Logger::Remote
-        http_adapter = Rager.config.http_adapter
-        url = Rager.config.url
-        api_key = Rager.config.api_key
-
-        unless url && api_key
-          raise Rager::Errors::CredentialsError.new("Rager Cloud", details: "Missing url or api_key for remote logging")
-        end
-
-        headers = {
-          "Content-Type" => "application/json",
-          "Authorization" => "Bearer #{api_key}"
-        }
+      Rager::Utils::Http.log_remote(
+        "results",
+        to_json,
+        !success?
+      )
+    end
 
-        request = Rager::Http::Request.new(
-          url: url,
-          verb: Rager::Http::Verb::Post,
-          headers: headers,
-          body: json
-        )
+    private
 
-        response = http_adapter.make_request(request)
+    sig { returns(T.nilable(Rager::Types::NonStreamOutput)) }
+    def display_output
+      return nil unless success?
+      return @consumed ? value : "[STREAM]" if stream?
 
-        unless response.success?
-          logger.error("Failed to log to remote server: #{response.status} #{response.body}")
-        end
-      end
+      T.cast(@output, Rager::Types::NonStreamOutput)
     end
   end
 end

data/lib/rager/search/options.rb

@@ -9,9 +9,10 @@ module Rager
       extend T::Sig
       include Rager::Options
 
-      const :provider, String, default: "jina"
-      const :max_length, T.nilable(Integer)
+      const :provider, Symbol, default: :jina
       const :n, T.nilable(Integer)
+      const :pagination, T.nilable(Integer)
+      const :max_length, T.nilable(Integer)
       const :api_key, T.nilable(String)
       const :timeout, T.nilable(Numeric)
     end

data/lib/rager/search/providers/jina.rb

@@ -22,18 +22,21 @@ module Rager
           api_key = options.api_key || ENV["JINA_API_KEY"]
           raise Rager::Errors::CredentialsError.new("Jina", env_var: ["JINA_API_KEY"]) if api_key.nil?
 
-          params = {"q" => query}
+          base_url = "https://s.jina.ai"
+          url = "#{base_url}/?#{URI.encode_www_form({"q" => query})}"
+          pagination = options.pagination
+          url = "#{url}&page=#{pagination}" if pagination && pagination > 1
 
           headers = {
             "Accept" => "application/json",
             "X-Return-Format" => "markdown"
-          }
-
-          headers["Authorization"] = "Bearer #{api_key}" if api_key
+          }.tap do |h|
+            h["Authorization"] = "Bearer #{api_key}" if api_key
+          end
 
           request = Rager::Http::Request.new(
             verb: Rager::Http::Verb::Get,
-            url: "https://s.jina.ai/?#{URI.encode_www_form(params)}",
+            url: url,
             headers: headers,
             timeout: options.timeout || Rager.config.timeout
           )
@@ -48,18 +51,16 @@ module Rager
 
           search_results = search_results.first(options.n) if options.n
 
-          urls = search_results.map { |r| r["url"] }
-          titles = search_results.map { |r| r["title"] }
-          contents = search_results.map { |r|
+          search_results.map do |r|
             content = r["content"] || r["description"]
-            options.max_length ? content&.slice(0, options.max_length) : content
-          }
+            content = content&.slice(0, options.max_length) if options.max_length
 
-          Rager::Search::Output.new(
-            urls: urls,
-            titles: titles,
-            contents: contents
-          )
+            Rager::Search::Result.new(
+              url: r["url"],
+              title: r["title"],
+              content: content
+            )
+          end
         end
       end
     end

data/lib/rager/search/result.rb

@@ -0,0 +1,21 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module Search
+    class Result < T::Struct
+      extend T::Sig
+
+      const :url, String
+      const :title, String
+      const :content, String
+
+      sig { params(options: T.untyped).returns(String) }
+      def to_json(options = nil)
+        serialize.to_json(options)
+      end
+    end
+  end
+end

data/lib/rager/template/input.rb

@@ -0,0 +1,20 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+
+module Rager
+  module Template
+    class Input < T::Struct
+      extend T::Sig
+
+      const :template, String
+      const :variables, T::Hash[Symbol, T.untyped]
+
+      sig { params(options: T.untyped).returns(String) }
+      def to_json(options = nil)
+        serialize.to_json(options)
+      end
+    end
+  end
+end

data/lib/rager/template/options.rb

@@ -7,7 +7,7 @@ module Rager
       include Rager::Options
       extend T::Sig
 
-      const :provider, String, default: "erb"
+      const :provider, Symbol, default: :erb
     end
   end
 end

data/lib/rager/template/providers/abstract.rb

@@ -14,12 +14,11 @@ module Rager
 
         sig do
           abstract.params(
-            template: String,
-            variables: T::Hash[Symbol, T.untyped],
+            input: Rager::Template::Input,
             options: Rager::Template::Options
           ).returns(Rager::Types::TemplateOutput)
         end
-        def template(template, variables, options)
+        def template(input, options)
         end
       end
     end

data/lib/rager/template/providers/erb.rb

@@ -11,15 +11,14 @@ module Rager
 
         sig do
           override.params(
-            template: String,
-            variables: T::Hash[Symbol, T.untyped],
+            input: Rager::Template::Input,
             options: Rager::Template::Options
           ).returns(Rager::Types::TemplateOutput)
         end
-        def template(template, variables, options)
-          ERB.new(template).result_with_hash(variables)
+        def template(input, options)
+          ERB.new(input.template).result_with_hash(input.variables)
         rescue SyntaxError, NameError => e
-          raise Rager::Errors::TemplateError.new(template, variables, details: e.message)
+          raise Rager::Errors::TemplateError.new(input.template, input.variables, details: e.message)
         end
       end
     end

data/lib/rager/template/providers/mustache.rb

@@ -0,0 +1,30 @@
+# typed: strict
+# frozen_string_literal: true
+
+module Rager
+  module Template
+    module Providers
+      class Mustache < Rager::Template::Providers::Abstract
+        extend T::Sig
+
+        sig do
+          override.params(
+            input: Rager::Template::Input,
+            options: Rager::Template::Options
+          ).returns(Rager::Types::TemplateOutput)
+        end
+        def template(input, options)
+          begin
+            require "mustache"
+          rescue LoadError
+            raise Rager::Errors::DependencyError.new("mustache", details: "Please install the mustache gem to use mustache templates")
+          end
+
+          ::Mustache.render(input.template, input.variables)
+        rescue ::Mustache::Parser::SyntaxError, ::Mustache::ContextMiss => e
+          raise Rager::Errors::TemplateError.new(input.template, input.variables, details: e.message)
+        end
+      end
+    end
+  end
+end

data/lib/rager/types.rb

@@ -1,71 +1,78 @@
 # typed: strict
 # frozen_string_literal: true
 
-require_relative "rerank/output"
-require_relative "search/output"
-
 module Rager
   module Types
     extend T::Sig
 
     ChatInput = T.type_alias { T::Array[Rager::Chat::Message] }
     EmbedInput = T.type_alias { T::Array[String] }
-    ImageGenInput = T.type_alias { String }
-    MeshGenInput = T.type_alias { String }
-    RerankInput = T.type_alias { {query: String, documents: T.any(T::Array[String], T::Array[T::Hash[Symbol, String]])} }
+    ImageInput = T.type_alias { String }
+    MeshInput = T.type_alias { String }
+    RerankInput = T.type_alias { Rager::Rerank::Input }
     SearchInput = T.type_alias { String }
-    TemplateInput = T.type_alias { {template: String, variables: T::Hash[Symbol, T.untyped]} }
+    TemplateInput = T.type_alias { Rager::Template::Input }
     Input = T.type_alias {
       T.any(
         ChatInput,
         EmbedInput,
-        ImageGenInput,
-        MeshGenInput,
+        ImageInput,
+        MeshInput,
         RerankInput,
         SearchInput,
-        TemplateInput,
-        Rager::Result
+        TemplateInput
       )
     }
 
-    ChatStream = T.type_alias { T::Enumerator[Rager::Chat::MessageDelta] }
-    Stream = T.type_alias { ChatStream }
+    ChatStreamOutput = T.type_alias { T::Enumerator[Rager::Chat::MessageDelta] }
+    Stream = T.type_alias { ChatStreamOutput }
 
     ChatBuffer = T.type_alias { T::Array[Rager::Chat::MessageDelta] }
     Buffer = T.type_alias { ChatBuffer }
 
-    ChatNonStream = T.type_alias { T.any(String, T::Array[String]) }
+    ChatNonStreamOutput = T.type_alias { T.any(String, T::Array[String]) }
 
-    ChatOutput = T.type_alias { T.any(ChatNonStream, ChatStream) }
+    ChatOutput = T.type_alias { T.any(ChatNonStreamOutput, ChatStreamOutput) }
     EmbedOutput = T.type_alias { T::Array[T::Array[Float]] }
-    ImageGenOutput = T.type_alias { String }
-    MeshGenOutput = T.type_alias { String }
-    RerankOutput = T.type_alias { Rager::Rerank::Output }
-    SearchOutput = T.type_alias { Rager::Search::Output }
+    ImageOutput = T.type_alias { String }
+    MeshOutput = T.type_alias { String }
+    RerankOutput = T.type_alias { T::Array[Rager::Rerank::Result] }
+    SearchOutput = T.type_alias { T::Array[Rager::Search::Result] }
     TemplateOutput = T.type_alias { String }
     Output = T.type_alias {
       T.any(
         ChatOutput,
         EmbedOutput,
-        ImageGenOutput,
-        MeshGenOutput,
+        ImageOutput,
+        MeshOutput,
         RerankOutput,
         SearchOutput,
         TemplateOutput
       )
     }
 
-    ChatNonStreamOutput = T.type_alias { ChatNonStream }
     NonStreamOutput = T.type_alias {
       T.any(
         ChatNonStreamOutput,
         EmbedOutput,
-        ImageGenOutput,
-        MeshGenOutput,
+        ImageOutput,
+        MeshOutput,
         RerankOutput,
         SearchOutput,
         TemplateOutput
       )
     }
+
+    Options = T.type_alias {
+      T.any(
+        Rager::Chat::Options,
+        Rager::Embed::Options,
+        Rager::Image::Options,
+        Rager::Mesh::Options,
+        Rager::Rerank::Options,
+        Rager::Search::Options,
+        Rager::Template::Options
+      )
+    }
   end
 end