cloudflare-ai 0.2.0 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 4a98e10b8688a9f4a0853534b104c9a05ceaa91e4405c71d4dde52fdd1953ce7
-  data.tar.gz: 5921d2c627b91dfcf58375b39e9d0781f839e880d8c996146d7002f9ba48f365
+  metadata.gz: b0a6b80e669a5c4d64ec87433ed1c3426ee0a2a1087016773c5648c799b34912
+  data.tar.gz: fe1d3b2afb5f70cbe476b1b8e006819b2e6e469184169993a15c7c4c34b89ca8
 SHA512:
-  metadata.gz: a2c1f4389bea67083ef6e19e12620b29850bdf818e4bbce1e3dcd4e5916175fabf69f04fe33b36ae539ec43f622aeb3456f3fff34ec9f236e745fe7f853cdbdc
-  data.tar.gz: c3e48f425b89449cae0bcbc480ee23c94d6494972f950070f3bb18e013711da95b45d7b6e3b2c2e7c9069ef458cc384ddf76ef3ea3a55cc61e46453bd190e189
+  metadata.gz: 84dcec88bcd44e4a7b004235b15b559d924c1e8dfce552234b2d9286762ab2a12d282954e8d863523b1c9170c24f49a772da7320b5b475b66b79b141ac2c532d
+  data.tar.gz: 0f676b2accb65edc63b2b9b1c2b8a532ebfedba588dfdcb1ae38ccfdf7edb4869c1ebee2060e887e5ac166772cda91d3e43e2ca650e878eb63b8c964d7fc19e1

data/README.md

@@ -1,35 +1,40 @@
-cloudflare-ai:  a Cloudflare Workers AI client for ruby developers
+Cloudflare Workers AI API client for ruby
 ---
 Cloudflare is testing its [Workers AI](https://blog.cloudflare.com/workers-ai) API. 
 Hopefully this project makes it easier for ruby-first developers to consume 
 Cloudflare's latest and greatest. 
 
-I'm a lawyer who codes, and I'm really interested in applying retrieval-augmented
+
+![Tests status](https://github.com/ajaynomics/cloudflare-ai/actions/workflows/ci.yml/badge.svg?branch=main)
+[![Gem Version](https://badge.fury.io/rb/cloudflare-ai.svg)](https://badge.fury.io/rb/cloudflare-ai)
+![GitHub License](https://img.shields.io/github/license/ajaynomics/cloudflare-ai)
+
+I'm really interested in applying retrieval-augmented
 generation to make legal services more accessible. [Email me](mailto:cloudflare-ai@krishnan.ca).
 
 If you're looking for legal help, it's best to book a slot via https://www.krishnan.ca.
 
-## Todo
+# Todo
 It's still early days, and here are my immediate priorities:
-* [ ] Support for streamed responses
-* [ ] CI pipeline and guidance for contributors
+* [x] Support for streamed responses
+* [x] CI pipeline
 * [ ] Support for more AI model categories
   * [x] [Text Generation](https://developers.cloudflare.com/workers-ai/models/text-generation/)
-  * [ ] [Text Embeddings](https://developers.cloudflare.com/workers-ai/models/text-embeddings/)
+  * [x] [Text Embeddings](https://developers.cloudflare.com/workers-ai/models/text-embeddings/)
   * [ ] [Text Classification](https://developers.cloudflare.com/workers-ai/models/text-classification/)
   * [ ] [Image Classification](https://developers.cloudflare.com/workers-ai/models/image-classification/)
   * [ ] [Translation](https://developers.cloudflare.com/workers-ai/models/translation/)
   * [ ] [Text-to-Image](https://developers.cloudflare.com/workers-ai/models/text-to-image/)
   * [ ] [Automatic Speech Recognition](https://developers.cloudflare.com/workers-ai/models/speech-recognition/)
 
-## Table of Contents
+# Table of Contents
 
 - [Installation](#installation)
 - [Usage](#usage)
 - [Logging](#logging)
 - [Development](#development)
 
-## Installation
+# Installation
 
 Install the gem and add to the application's Gemfile by executing:
 
@@ -39,26 +44,51 @@ If bundler is not being used to manage dependencies, install the gem by executin
 
     gem install cloudflare-ai
 
-## Usage
+# Usage
 
 ```ruby
 require "cloudflare/ai"
 ```
 
-### Cloudflare Workers AI
+## Cloudflare Workers AI
 Please visit the [Cloudflare Workers AI website](https://developers.cloudflare.com/workers-ai/) for more details.
 Thiis gem provides a client that wraps around [Cloudflare's REST API](https://developers.cloudflare.com/workers-ai/get-started/rest-api/).
 
 
-### Client
+## Client
 
 ```ruby
 client = Cloudflare::AI::Client.new(account_id: ENV["CLOUDFLARE_ACCOUNT_ID"], api_token: ENV["CLOUDFLARE_API_TOKEN"])
 ```
 
+
+### Text generation (chat / scoped prompt)
+```ruby
+messages = [
+  Cloudflare::AI::Message.new(role: "system", content: "You are a big fan of Cloudflare and Ruby."),
+  Cloudflare::AI::Message.new(role: "user", content: "What is your favourite tech stack?"),
+  Cloudflare::AI::Message.new(role: "assistant", content: "I love building with Ruby on Rails and Cloudflare!"),
+  Cloudflare::AI::Message.new(role: "user", content: "Really? You like Cloudflare even though there isn't great support for Ruby?"),
+]
+result = client.chat(messages: messages)
+puts result.response # => "Yes, I love Cloudflare!"
+```
+
+#### Streaming responses
+Responses will be streamed back to the client using Server Side Events (SSE) if a block is passed to the `chat` or `complete` method.
+```ruby
+result = client.complete(prompt: "Hi!") { |data| puts data}
+# {"response":" "}
+# {"response":" Hello"}
+# {"response":" there"}
+# {"response":"!"}
+# {"response":""}
+# [DONE]
+
+```
 #### Result object
-All invocations of the client return a `Cloudflare::AI::Result` object. This object's serializable JSON output is 
-based on the raw response from the Cloudflare API. 
+All invocations of the `prompt` and `chat` methods return a `Cloudflare::AI::Results::TextGeneration` object. This object's serializable JSON output is
+based on the raw response from the Cloudflare API.
 
 ```ruby
 result = client.complete(prompt: "What is your name?")
@@ -77,30 +107,22 @@ puts result.to_json # => {"result":null,"success":false,"errors":[{"code":7009,"
 ```
 
 
-#### Text generation (chat / scoped prompt)
+### Text embedding
 ```ruby
-messages = [
-  Cloudflare::AI::Message.new(role: "system", content: "You are a big fan of Cloudflare and Ruby."),
-  Cloudflare::AI::Message.new(role: "user", content: "What is your favourite tech stack?"),
-  Cloudflare::AI::Message.new(role: "assistant", content: "I love building with Ruby on Rails and Cloudflare!"),
-  Cloudflare::AI::Message.new(role: "user", content: "Really? You like Cloudflare even though there isn't great support for Ruby?"),
-]
-result = client.chat(messages: messages)
-puts result.response # => "Yes, I love Cloudflare!"
+result = client.embed(text: "Hello")
+p result.shape # => [1, 768] # (1 embedding, 768 dimensions per embedding)
+p result.embedding # => [[-0.008496830239892006, 0.001376907923258841, -0.0323275662958622, ...]]
 ```
 
-##### Streaming responses
-Responses will be streamed back to the client using Server Side Events (SSE) if a block is passed to the `chat` or `complete` method.
+The input can be either a string (as above) or an array of strings:
 ```ruby
-result = client.complete(prompt: "Hi!") { |data| puts data}
-# {"response":" "}
-# {"response":" Hello"}
-# {"response":" there"}
-# {"response":"!"}
-# {"response":""}
-# [DONE]
+result = client.embed(text: ["Hello", "World"])
 ```
-## Logging
+
+#### Result object
+All invocations of the `embedding` methods return a `Cloudflare::AI::Results::TextEmbedding`. 
+
+# Logging
 
 This gem uses standard logging mechanisms and defaults to `:warn` level. Most messages are at info level, but we will add debug or warn statements as needed.
 To show all log messages:
@@ -113,12 +135,12 @@ You can use this logger as you would the default ruby logger. For example:
 ```ruby
 Cloudflare::AI.logger = Logger.new($stdout)
 ```
-## Development
+# Development
 
 1. `git clone https://github.com/ajaynomics/cloudflare-ai.git`
 2. `bundle exec rake` to ensure that the tests pass and to run standardrb
 
-## Contributing
+# Contributing
 
 Bug reports and pull requests are welcome on GitHub at https://github.com/ajaynomics/cloudflare-ai.
 

data/lib/cloudflare/ai/client.rb

@@ -2,6 +2,8 @@ require "event_stream_parser"
 require "faraday"
 
 class Cloudflare::AI::Client
+  include Cloudflare::AI::Clients::TextGenerationHelpers
+
   attr_reader :url, :account_id, :api_token
 
   def initialize(account_id:, api_token:)
@@ -9,50 +11,28 @@ class Cloudflare::AI::Client
     @api_token = api_token
   end
 
-  def complete(prompt:, model_name: models[:text_generation].first, &block)
+  def chat(messages:, model_name: default_text_generation_model_name, &block)
     url = service_url_for(account_id: account_id, model_name: model_name)
     stream = block ? true : false
-    payload = create_payload({prompt: prompt}, stream: stream)
-    post_request(url, payload, &block)
+    payload = create_streamable_payload({messages: messages.map(&:serializable_hash)}, stream: stream)
+    post_streamable_request(url, payload, &block)
   end
 
-  def chat(messages:, model_name: models[:text_generation].first, &block)
+  def complete(prompt:, model_name: default_text_generation_model_name, &block)
     url = service_url_for(account_id: account_id, model_name: model_name)
     stream = block ? true : false
-    payload = create_payload({messages: messages.map(&:serializable_hash)}, stream: stream)
-    post_request(url, payload, &block)
-  end
-
-  def models
-    {
-      text_generation: %w[@cf/meta/llama-2-7b-chat-fp16 @cf/meta/llama-2-7b-chat-int8 @cf/mistral/mistral-7b-instruct-v0.1 @hf/thebloke/codellama-7b-instruct-awq],
-      speech_recognition: %w[@cf/openai/whisper],
-      translation: %w[@cf/meta/m2m100-1.2b],
-      text_classification: %w[@cf/huggingface/distilbert-sst-2-int8],
-      image_classification: %w[@cf/huggingface/distilbert-sst-2-int8],
-      text_to_image: %w[@cf/stabilityai/stable-diffusion-xl-base-1.0],
-      text_embeddings: %w[@cf/baai/bge-base-en-v1.5 @cf/baai/bge-large-en-v1.5 @cf/baai/bge-small-en-v1.5]
-    }.freeze
+    payload = create_streamable_payload({prompt: prompt}, stream: stream)
+    post_streamable_request(url, payload, &block)
   end
 
-  private
+  def embed(text:, model_name: Cloudflare::AI::Models.text_embedding.first)
+    url = service_url_for(account_id: account_id, model_name: model_name)
+    payload = {text: text}.to_json
 
-  def create_payload(data, stream: false)
-    data.merge({stream: stream}).to_json
+    Cloudflare::AI::Results::TextEmbedding.new(connection.post(url, payload).body)
   end
 
-  def post_request(url, payload, &block)
-    if block
-      parser = EventStreamParser::Parser.new
-      connection.post(url, payload) do |response|
-        response.options.on_data = parser.stream do |_type, data, _id, _reconnection_time, _size|
-          yield data
-        end
-      end
-    else
-      Cloudflare::AI::Result.new(connection.post(url, payload).body)
-    end
-  end
+  private
 
   def connection
     @connection ||= ::Faraday.new(headers: {Authorization: "Bearer #{api_token}"})

data/lib/cloudflare/ai/clients/text_generation_helpers.rb

@@ -0,0 +1,28 @@
+module Cloudflare
+  module AI
+    module Clients
+      module TextGenerationHelpers
+        def default_text_generation_model_name
+          Cloudflare::AI::Models.text_generation.first
+        end
+
+        def create_streamable_payload(data, stream: false)
+          data.merge({stream: stream}).to_json
+        end
+
+        def post_streamable_request(url, payload, &block)
+          if block
+            parser = EventStreamParser::Parser.new
+            connection.post(url, payload) do |response|
+              response.options.on_data = parser.stream do |_type, data, _id, _reconnection_time, _size|
+                yield data
+              end
+            end
+          else
+            Cloudflare::AI::Results::TextGeneration.new(connection.post(url, payload).body)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/cloudflare/ai/models.rb

@@ -0,0 +1,43 @@
+class Cloudflare::AI::Models
+  class << self
+    def text_generation
+      %w[@cf/meta/llama-2-7b-chat-fp16 @cf/meta/llama-2-7b-chat-int8 @cf/mistral/mistral-7b-instruct-v0.1 @hf/thebloke/codellama-7b-instruct-awq]
+    end
+
+    def speech_recognition
+      %w[@cf/openai/whisper]
+    end
+
+    def translation
+      %w[@cf/meta/m2m100-1.2b]
+    end
+
+    def text_classification
+      %w[@cf/huggingface/distilbert-sst-2-int8]
+    end
+
+    def image_classification
+      %w[@cf/huggingface/distilbert-sst-2-int8]
+    end
+
+    def text_to_image
+      %w[@cf/stabilityai/stable-diffusion-xl-base-1.0]
+    end
+
+    def text_embedding
+      %w[@cf/baai/bge-base-en-v1.5 @cf/baai/bge-large-en-v1.5 @cf/baai/bge-small-en-v1.5]
+    end
+
+    def all
+      {
+        text_generation: text_generation,
+        speech_recognition: speech_recognition,
+        translation: translation,
+        text_classification: text_classification,
+        image_classification: image_classification,
+        text_to_image: text_to_image,
+        text_embeddings: text_embedding
+      }
+    end
+  end
+end

data/lib/cloudflare/ai/result.rb

@@ -1,32 +1,45 @@
 require "active_support/core_ext/hash/indifferent_access"
+require "json"
 
 class Cloudflare::AI::Result
-  attr_reader :result, :success, :errors, :messages
-
-  def initialize(json)
-    @json = json
-    @json = JSON.parse(@json) unless @json.is_a?(Hash)
-    @json = @json.with_indifferent_access
+  def initialize(json_string_or_ruby_hash)
+    @result_data = parse_data(json_string_or_ruby_hash)
+  end
 
-    @result = @json["result"]
-    @success = @json["success"]
-    @errors = @json["errors"]
-    @messages = @json["messages"]
+  def result
+    result_data[:result]
   end
 
-  def to_json
-    @json.to_json
+  def success?
+    success == true
   end
 
   def failure?
     !success?
   end
 
-  def success?
-    success
+  def errors
+    result_data.dig(:errors)
+  end
+
+  def messages
+    result_data.dig(:messages)
+  end
+
+  def to_json
+    result_data.to_json
+  end
+
+  private
+
+  attr_reader :result_data
+
+  def success
+    result_data[:success]
   end
 
-  def response
-    result.with_indifferent_access["response"]
+  def parse_data(input)
+    input = JSON.parse(input) if input.is_a?(String)
+    input.with_indifferent_access
   end
 end

data/lib/cloudflare/ai/results/text_embedding.rb

@@ -0,0 +1,9 @@
+class Cloudflare::AI::Results::TextEmbedding < Cloudflare::AI::Result
+  def shape
+    result&.dig(:shape) # nil if no shape
+  end
+
+  def data
+    result&.dig(:data) # nil if no data
+  end
+end

data/lib/cloudflare/ai/results/text_generation.rb

@@ -0,0 +1,5 @@
+class Cloudflare::AI::Results::TextGeneration < Cloudflare::AI::Result
+  def response
+    result&.dig(:response) # nil if no response
+  end
+end

data/lib/cloudflare/ai/version.rb

@@ -2,6 +2,6 @@
 
 module Cloudflare
   module AI
-    VERSION = "0.2.0"
+    VERSION = "0.3.0"
   end
 end

data/lib/cloudflare/ai.rb

@@ -1,5 +1,3 @@
-# frozen_string_literal: true
-
 require_relative "ai/version"
 require "logger"
 require "pathname"
@@ -27,6 +25,5 @@ module Cloudflare
     @root = Pathname.new(__dir__)
 
     class Error < StandardError; end
-    # Your code goes here...
   end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: cloudflare-ai
 version: !ruby/object:Gem::Version
-  version: 0.2.0
+  version: 0.3.0
 platform: ruby
 authors:
 - Ajay Krishnan
@@ -93,9 +93,13 @@ files:
 - README.md
 - lib/cloudflare/ai.rb
 - lib/cloudflare/ai/client.rb
+- lib/cloudflare/ai/clients/text_generation_helpers.rb
 - lib/cloudflare/ai/contextual_logger.rb
 - lib/cloudflare/ai/message.rb
+- lib/cloudflare/ai/models.rb
 - lib/cloudflare/ai/result.rb
+- lib/cloudflare/ai/results/text_embedding.rb
+- lib/cloudflare/ai/results/text_generation.rb
 - lib/cloudflare/ai/version.rb
 homepage: https://rubygems.org/gems/cloudflare-ai
 licenses: