llm.rb 0.1.0

data/lib/llm/providers/anthropic/response_parser.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+class LLM::Anthropic
+  module ResponseParser
+    def parse_embedding(body)
+      {
+        model: body["model"],
+        embeddings: body["data"].map { _1["embedding"] },
+        total_tokens: body.dig("usage", "total_tokens")
+      }
+    end
+
+    ##
+    # @param [Hash] body
+    #  The response body from the LLM provider
+    # @return [Hash]
+    def parse_completion(body)
+      {
+        model: body["model"],
+        choices: body["content"].map do
+          # TODO: don't hardcode role
+          LLM::Message.new("assistant", _1["text"], {completion: self})
+        end,
+        prompt_tokens: body.dig("usage", "input_tokens"),
+        completion_tokens: body.dig("usage", "output_tokens")
+      }
+    end
+  end
+end

data/lib/llm/providers/anthropic.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module LLM
+  ##
+  # The Anthropic class implements a provider for
+  # [Anthropic](https://www.anthropic.com)
+  class Anthropic < Provider
+    require_relative "anthropic/error_handler"
+    require_relative "anthropic/response_parser"
+    require_relative "anthropic/format"
+    include Format
+
+    HOST = "api.anthropic.com"
+    DEFAULT_PARAMS = {max_tokens: 1024, model: "claude-3-5-sonnet-20240620"}.freeze
+
+    ##
+    # @param secret (see LLM::Provider#initialize)
+    def initialize(secret, **)
+      super(secret, host: HOST, **)
+    end
+
+    ##
+    # @param input (see LLM::Provider#embed)
+    # @return (see LLM::Provider#embed)
+    def embed(input, **params)
+      req = Net::HTTP::Post.new ["api.voyageai.com/v1", "embeddings"].join("/")
+      body = {input:, model: "voyage-2"}.merge!(params)
+      req = preflight(req, body)
+      res = request(@http, req)
+      Response::Embedding.new(res).extend(response_parser)
+    end
+
+    ##
+    # @see https://docs.anthropic.com/en/api/messages Anthropic docs
+    # @param prompt (see LLM::Provider#complete)
+    # @param role (see LLM::Provider#complete)
+    # @return (see LLM::Provider#complete)
+    def complete(prompt, role = :user, **params)
+      req = Net::HTTP::Post.new ["/v1", "messages"].join("/")
+      messages = [*(params.delete(:messages) || []), Message.new(role, prompt)]
+      params = DEFAULT_PARAMS.merge(params)
+      body = {messages: format(messages)}.merge!(params)
+      req = preflight(req, body)
+      res = request(@http, req)
+      Response::Completion.new(res).extend(response_parser)
+    end
+
+    private
+
+    def auth(req)
+      req["anthropic-version"] = "2023-06-01"
+      req["x-api-key"] = @secret
+    end
+
+    def response_parser
+      LLM::Anthropic::ResponseParser
+    end
+
+    def error_handler
+      LLM::Anthropic::ErrorHandler
+    end
+  end
+end

data/lib/llm/providers/gemini/error_handler.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+class LLM::Gemini
+  class ErrorHandler
+    ##
+    # @return [Net::HTTPResponse]
+    #  Non-2XX response from the server
+    attr_reader :res
+
+    ##
+    # @param [Net::HTTPResponse] res
+    #  The response from the server
+    # @return [LLM::Gemini::ErrorHandler]
+    def initialize(res)
+      @res = res
+    end
+
+    ##
+    # @raise [LLM::Error]
+    #  Raises a subclass of {LLM::Error LLM::Error}
+    def raise_error!
+      case res
+      when Net::HTTPBadRequest
+        reason = body.dig("error", "details", 0, "reason")
+        if reason == "API_KEY_INVALID"
+          raise LLM::Error::Unauthorized.new { _1.response = res }, "Authentication error"
+        else
+          raise LLM::Error::BadResponse.new { _1.response = res }, "Unexpected response"
+        end
+      when Net::HTTPTooManyRequests
+        raise LLM::Error::RateLimit.new { _1.response = res }, "Too many requests"
+      else
+        raise LLM::Error::BadResponse.new { _1.response = res }, "Unexpected response"
+      end
+    end
+
+    private
+
+    def body
+      @body ||= JSON.parse(res.body)
+    end
+  end
+end

data/lib/llm/providers/gemini/format.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+class LLM::Gemini
+  module Format
+    ##
+    # @param [Array<LLM::Message>] messages
+    #  The messages to format
+    # @return [Array<Hash>]
+    def format(messages)
+      messages.map { {role: _1.role, parts: [format_content(_1.content)]} }
+    end
+
+    private
+
+    ##
+    # @param [String, LLM::File] content
+    #  The content to format
+    # @return [String, Hash]
+    #  The formatted content
+    def format_content(content)
+      if LLM::File === content
+        file = content
+        {
+          inline_data: {mime_type: file.mime_type, data: [File.binread(file.path)].pack("m0")}
+        }
+      else
+        {text: content}
+      end
+    end
+  end
+end

data/lib/llm/providers/gemini/response_parser.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+class LLM::Gemini
+  module ResponseParser
+    def parse_embedding(body)
+      {
+        model: "text-embedding-004",
+        embeddings: body.dig("embedding", "values")
+      }
+    end
+
+    ##
+    # @param [Hash] body
+    #  The response body from the LLM provider
+    # @return [Hash]
+    def parse_completion(body)
+      {
+        model: body["modelVersion"],
+        choices: body["candidates"].map do
+          LLM::Message.new(
+            _1.dig("content", "role"),
+            _1.dig("content", "parts", 0, "text"),
+            {completion: self}
+          )
+        end,
+        prompt_tokens: body.dig("usageMetadata", "promptTokenCount"),
+        completion_tokens: body.dig("usageMetadata", "candidatesTokenCount")
+      }
+    end
+  end
+end

data/lib/llm/providers/gemini.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+module LLM
+  ##
+  # The Gemini class implements a provider for
+  # [Gemini](https://ai.google.dev/)
+  class Gemini < Provider
+    require_relative "gemini/error_handler"
+    require_relative "gemini/response_parser"
+    require_relative "gemini/format"
+    include Format
+
+    HOST = "generativelanguage.googleapis.com"
+    DEFAULT_PARAMS = {model: "gemini-1.5-flash"}.freeze
+
+    ##
+    # @param secret (see LLM::Provider#initialize)
+    def initialize(secret, **)
+      super(secret, host: HOST, **)
+    end
+
+    ##
+    # @param input (see LLM::Provider#embed)
+    # @return (see LLM::Provider#embed)
+    def embed(input, **params)
+      path = ["/v1beta/models", "text-embedding-004"].join("/")
+      req = Net::HTTP::Post.new [path, "embedContent"].join(":")
+      body = {content: {parts: [{text: input}]}}
+      req = preflight(req, body)
+      res = request @http, req
+      Response::Embedding.new(res).extend(response_parser)
+    end
+
+    ##
+    # @see https://ai.google.dev/api/generate-content#v1beta.models.generateContent Gemini docs
+    # @param prompt (see LLM::Provider#complete)
+    # @param role (see LLM::Provider#complete)
+    # @return (see LLM::Provider#complete)
+    def complete(prompt, role = :user, **params)
+      params = DEFAULT_PARAMS.merge(params)
+      path = ["/v1beta/models", params.delete(:model)].join("/")
+      req = Net::HTTP::Post.new [path, "generateContent"].join(":")
+      messages = [*(params.delete(:messages) || []), LLM::Message.new(role, prompt)]
+      body = {contents: format(messages)}
+      req = preflight(req, body)
+      res = request(@http, req)
+      Response::Completion.new(res).extend(response_parser)
+    end
+
+    private
+
+    def auth(req)
+      req.path.replace [req.path, URI.encode_www_form(key: @secret)].join("?")
+    end
+
+    def response_parser
+      LLM::Gemini::ResponseParser
+    end
+
+    def error_handler
+      LLM::Gemini::ErrorHandler
+    end
+  end
+end

data/lib/llm/providers/ollama/error_handler.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+class LLM::Ollama
+  class ErrorHandler
+    ##
+    # @return [Net::HTTPResponse]
+    #  Non-2XX response from the server
+    attr_reader :res
+
+    ##
+    # @param [Net::HTTPResponse] res
+    #  The response from the server
+    # @return [LLM::OpenAI::ErrorHandler]
+    def initialize(res)
+      @res = res
+    end
+
+    ##
+    # @raise [LLM::Error]
+    #  Raises a subclass of {LLM::Error LLM::Error}
+    def raise_error!
+      case res
+      when Net::HTTPUnauthorized
+        raise LLM::Error::Unauthorized.new { _1.response = res }, "Authentication error"
+      when Net::HTTPTooManyRequests
+        raise LLM::Error::RateLimit.new { _1.response = res }, "Too many requests"
+      else
+        raise LLM::Error::BadResponse.new { _1.response = res }, "Unexpected response"
+      end
+    end
+  end
+end

data/lib/llm/providers/ollama/format.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+class LLM::Ollama
+  module Format
+    ##
+    # @param [Array<LLM::Message>] messages
+    #  The messages to format
+    # @return [Array<Hash>]
+    def format(messages)
+      messages.map { {role: _1.role, content: format_content(_1.content)} }
+    end
+
+    private
+
+    ##
+    # @param [String, URI] content
+    #  The content to format
+    # @return [String, Hash]
+    #  The formatted content
+    def format_content(content)
+      if URI === content
+        [{type: :image_url, image_url: {url: content.to_s}}]
+      else
+        content
+      end
+    end
+  end
+end

data/lib/llm/providers/ollama/response_parser.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+class LLM::Ollama
+  module ResponseParser
+    ##
+    # @param [Hash] body
+    #  The response body from the LLM provider
+    # @return [Hash]
+    def parse_completion(body)
+      {
+        model: body["model"],
+        choices: [LLM::Message.new(*body["message"].values_at("role", "content"), {completion: self})],
+        prompt_tokens: body.dig("prompt_eval_count"),
+        completion_tokens: body.dig("eval_count")
+      }
+    end
+  end
+end

data/lib/llm/providers/ollama.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module LLM
+  ##
+  # The Ollama class implements a provider for
+  # [Ollama](https://ollama.ai/)
+  class Ollama < Provider
+    require_relative "ollama/error_handler"
+    require_relative "ollama/response_parser"
+    require_relative "ollama/format"
+    include Format
+
+    HOST = "localhost"
+    DEFAULT_PARAMS = {model: "llama3.2", stream: false}.freeze
+
+    ##
+    # @param secret (see LLM::Provider#initialize)
+    def initialize(secret, **)
+      super(secret, host: HOST, port: 11434, ssl: false, **)
+    end
+
+    ##
+    # @see https://github.com/ollama/ollama/blob/main/docs/api.md#generate-a-chat-completion Ollama docs
+    # @param prompt (see LLM::Provider#complete)
+    # @param role (see LLM::Provider#complete)
+    # @return (see LLM::Provider#complete)
+    def complete(prompt, role = :user, **params)
+      req = Net::HTTP::Post.new ["/api", "chat"].join("/")
+      messages = [*(params.delete(:messages) || []), LLM::Message.new(role, prompt)]
+      params = DEFAULT_PARAMS.merge(params)
+      body = {messages: messages.map(&:to_h)}.merge!(params)
+      req = preflight(req, body)
+      res = request(@http, req)
+      Response::Completion.new(res).extend(response_parser)
+    end
+
+    private
+
+    def auth(req)
+      req["Authorization"] = "Bearer #{@secret}"
+    end
+
+    def response_parser
+      LLM::Ollama::ResponseParser
+    end
+
+    def error_handler
+      LLM::Ollama::ErrorHandler
+    end
+  end
+end

data/lib/llm/providers/openai/error_handler.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+class LLM::OpenAI
+  class ErrorHandler
+    ##
+    # @return [Net::HTTPResponse]
+    #  Non-2XX response from the server
+    attr_reader :res
+
+    ##
+    # @param [Net::HTTPResponse] res
+    #  The response from the server
+    # @return [LLM::OpenAI::ErrorHandler]
+    def initialize(res)
+      @res = res
+    end
+
+    ##
+    # @raise [LLM::Error]
+    #  Raises a subclass of {LLM::Error LLM::Error}
+    def raise_error!
+      case res
+      when Net::HTTPUnauthorized
+        raise LLM::Error::Unauthorized.new { _1.response = res }, "Authentication error"
+      when Net::HTTPTooManyRequests
+        raise LLM::Error::RateLimit.new { _1.response = res }, "Too many requests"
+      else
+        raise LLM::Error::BadResponse.new { _1.response = res }, "Unexpected response"
+      end
+    end
+  end
+end

data/lib/llm/providers/openai/format.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+class LLM::OpenAI
+  module Format
+    ##
+    # @param [Array<LLM::Message>] messages
+    #  The messages to format
+    # @return [Array<Hash>]
+    def format(messages)
+      messages.map { {role: _1.role, content: format_content(_1.content)} }
+    end
+
+    private
+
+    ##
+    # @param [String, URI] content
+    #  The content to format
+    # @return [String, Hash]
+    #  The formatted content
+    def format_content(content)
+      if URI === content
+        [{type: :image_url, image_url: {url: content.to_s}}]
+      else
+        content
+      end
+    end
+  end
+end

data/lib/llm/providers/openai/response_parser.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+class LLM::OpenAI
+  module ResponseParser
+    def parse_embedding(body)
+      {
+        model: body["model"],
+        embeddings: body.dig("data").map do |data|
+          data["embedding"]
+        end,
+        prompt_tokens: body.dig("usage", "prompt_tokens"),
+        total_tokens: body.dig("usage", "total_tokens")
+      }
+    end
+
+    ##
+    # @param [Hash] body
+    #  The response body from the LLM provider
+    # @return [Hash]
+    def parse_completion(body)
+      {
+        model: body["model"],
+        choices: body["choices"].map do
+          mesg = _1["message"]
+          logprobs = _1["logprobs"]
+          role, content = mesg.values_at("role", "content")
+          LLM::Message.new(role, content, {completion: self, logprobs:})
+        end,
+        prompt_tokens: body.dig("usage", "prompt_tokens"),
+        completion_tokens: body.dig("usage", "completion_tokens"),
+        total_tokens: body.dig("usage", "total_tokens")
+      }
+    end
+  end
+end

data/lib/llm/providers/openai.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+module LLM
+  ##
+  # The OpenAI class implements a provider for
+  # [OpenAI](https://platform.openai.com/)
+  class OpenAI < Provider
+    require_relative "openai/error_handler"
+    require_relative "openai/response_parser"
+    require_relative "openai/format"
+    include Format
+
+    HOST = "api.openai.com"
+    DEFAULT_PARAMS = {model: "gpt-4o-mini"}.freeze
+
+    ##
+    # @param secret (see LLM::Provider#initialize)
+    def initialize(secret, **)
+      super(secret, host: HOST, **)
+    end
+
+    ##
+    # @param input (see LLM::Provider#embed)
+    # @return (see LLM::Provider#embed)
+    def embed(input, **params)
+      req = Net::HTTP::Post.new ["/v1", "embeddings"].join("/")
+      body = {input:, model: "text-embedding-3-small"}.merge!(params)
+      req = preflight(req, body)
+      res = request @http, req
+      Response::Embedding.new(res).extend(response_parser)
+    end
+
+    ##
+    # @see https://platform.openai.com/docs/api-reference/chat/create OpenAI docs
+    # @param prompt (see LLM::Provider#complete)
+    # @param role (see LLM::Provider#complete)
+    # @return (see LLM::Provider#complete)
+    def complete(prompt, role = :user, **params)
+      req = Net::HTTP::Post.new ["/v1", "chat", "completions"].join("/")
+      messages = [*(params.delete(:messages) || []), Message.new(role, prompt)]
+      params = DEFAULT_PARAMS.merge(params)
+      body = {messages: format(messages)}.merge!(params)
+      req = preflight(req, body)
+      res = request(@http, req)
+      Response::Completion.new(res).extend(response_parser)
+    end
+
+    private
+
+    def auth(req)
+      req["Authorization"] = "Bearer #{@secret}"
+    end
+
+    def response_parser
+      LLM::OpenAI::ResponseParser
+    end
+
+    def error_handler
+      LLM::OpenAI::ErrorHandler
+    end
+  end
+end

data/lib/llm/response/completion.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module LLM
+  class Response::Completion < Response
+    ##
+    # @return [String]
+    #   Returns the model name used for the completion
+    def model
+      parsed[:model]
+    end
+
+    ##
+    # @return [Array<LLM::Message>]
+    #   Returns an array of messages
+    def choices
+      parsed[:choices]
+    end
+
+    ##
+    # @return [Integer]
+    #   Returns the count of prompt tokens
+    def prompt_tokens
+      parsed[:prompt_tokens]
+    end
+
+    ##
+    # @return [Integer]
+    #   Returns the count of completion tokens
+    def completion_tokens
+      parsed[:completion_tokens]
+    end
+
+    ##
+    # @return [Integer]
+    #   Returns the total count of tokens
+    def total_tokens
+      prompt_tokens + completion_tokens
+    end
+
+    private
+
+    ##
+    # @private
+    # @return [Hash]
+    #   Returns the parsed completion response from the provider
+    def parsed
+      @parsed ||= parse_completion(body)
+    end
+  end
+end

data/lib/llm/response/embedding.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module LLM
+  class Response::Embedding < Response
+    def model
+      parsed[:model]
+    end
+
+    def embeddings
+      parsed[:embeddings]
+    end
+
+    def total_tokens
+      parsed[:total_tokens]
+    end
+
+    private
+
+    def parsed
+      @parsed ||= parse_embedding(body)
+    end
+  end
+end

data/lib/llm/response.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module LLM
+  class Response
+    require "json"
+    require_relative "response/completion"
+    require_relative "response/embedding"
+
+    ##
+    # @return [Hash]
+    #  Returns the response body
+    attr_reader :body
+
+    ##
+    # @param [Net::HTTPResponse] res
+    #  HTTP response
+    # @return [LLM::Response]
+    #  Returns an instance of LLM::Response
+    def initialize(res)
+      @res = res
+      @body = JSON.parse(res.body)
+    end
+  end
+end

data/lib/llm/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module LLM
+  VERSION = "0.1.0"
+end