Add file search

Author: alexrudall

README.md

@@ -58,16 +58,6 @@ The engine options are currently "ada", "babbage", "curie" and "davinci". Hit th
     => [", there lived a great"]
 ```
 
-### Search
-
-Pass documents and a query string to get semantic search scores against each document:
-
-```
-    response = client.search(engine: "ada", documents: %w[washington hospital school], query: "president")
-    puts response["data"].map { |d| d["score"] }
-    => [202.0, 48.052, 19.247]
-```
-
 ### Files
 
 Put your data in a `.jsonl` file like this:
@@ -86,6 +76,22 @@ and pass the path to `client.files.upload` to upload it to OpenAI, and then inte
     client.files.delete(id: 123)
 ```
 
+### Search
+
+Pass documents and a query string to get semantic search scores against each document:
+
+```
+    response = client.search(engine: "ada", documents: %w[washington hospital school], query: "president")
+    puts response["data"].map { |d| d["score"] }
+    => [202.0, 48.052, 19.247]
+```
+
+You can alternatively search using the ID of a file you've uploaded:
+
+```
+    client.search(engine: "ada", file: "abc123", query: "happy")
+```
+
 ## Development
 
 After checking out the repo, run `bin/setup` to install dependencies. Then, run `rake spec` to run the tests. You can also run `bin/console` for an interactive prompt that will allow you to experiment.

lib/ruby/openai/client.rb

@@ -22,16 +22,14 @@ def files
       @files ||= OpenAI::Files.new(access_token: @access_token)
     end
 
-    def search(engine:, documents:, query:, version: default_version)
+    def search(engine:, query:, documents: nil, file: nil, version: default_version)
       self.class.post(
         "/#{version}/engines/#{engine}/search",
         headers: {
           "Content-Type" => "application/json",
           "Authorization" => "Bearer #{@access_token}"
         },
-        body: {
-          documents: documents, query: query
-        }.to_json
+        body: { file: file, query: query }.merge(documents ? { documents: documents } : { file: file }).to_json
       )
     end
 

spec/fixtures/cassettes/davinci_search_happy.yml

@@ -1,58 +1,85 @@
 RSpec.describe OpenAI::Client do
-  describe "#search" do
-    context "with a query" do
-      let(:query) { "the president" }
+  describe "#search", :vcr do
+    let(:cassette) { "#{engine} search #{query}".downcase }
 
-      context "with documents", :vcr do
-        let(:documents) { %w[washington hospital school] }
+    context "with a file" do
+      let(:filename) { "puppy.jsonl" }
+      let(:file) { File.join(RSPEC_ROOT, "fixtures/files", filename) }
+      let!(:file_id) do
+        response = VCR.use_cassette("files upload search") do
+          OpenAI::Client.new.files.upload(parameters: { file: file, purpose: "search" })
+        end
+        JSON.parse(response.body)["id"]
+      end
 
+      context "with engine: davinci" do
+        let(:query) { "happy" }
+        let(:engine) { "davinci" }
         let(:response) do
           OpenAI::Client.new.search(
             engine: engine,
-            documents: documents,
+            file: file_id,
             query: query
           )
         end
-        let(:best_match) { JSON.parse(response.body)["data"].max_by { |d| d["score"] }["document"] }
-        let(:cassette) { "#{engine} search #{query}".downcase }
 
-        context "with engine: ada" do
-          let(:engine) { "ada" }
+        it "finds the best match" do
+          VCR.use_cassette(cassette) do
+            expect(response.parsed_response['data'][0]["text"]).to eq("puppy A is happy")
+          end
+        end
+      end
+    end
+
+    context "with documents" do
+      let(:documents) { %w[washington hospital school] }
+      let(:query) { "the president" }
+      let(:response) do
+        OpenAI::Client.new.search(
+          engine: engine,
+          documents: documents,
+          query: query
+        )
+      end
+      let(:best_match) { JSON.parse(response.body)["data"].max_by { |d| d["score"] }["document"] }
+      let(:cassette) { "#{engine} search #{query}".downcase }
+
+      context "with engine: ada" do
+        let(:engine) { "ada" }
 
-          it "finds the best match" do
-            VCR.use_cassette(cassette) do
-              expect(documents[best_match]).to eq("washington")
-            end
+        it "finds the best match" do
+          VCR.use_cassette(cassette) do
+            expect(documents[best_match]).to eq("washington")
           end
         end
+      end
 
-        context "with engine: babbage" do
-          let(:engine) { "babbage" }
+      context "with engine: babbage" do
+        let(:engine) { "babbage" }
 
-          it "finds the best match" do
-            VCR.use_cassette(cassette) do
-              expect(documents[best_match]).to eq("washington")
-            end
+        it "finds the best match" do
+          VCR.use_cassette(cassette) do
+            expect(documents[best_match]).to eq("washington")
           end
         end
+      end
 
-        context "with engine: curie" do
-          let(:engine) { "curie" }
+      context "with engine: curie" do
+        let(:engine) { "curie" }
 
-          it "finds the best match" do
-            VCR.use_cassette(cassette) do
-              expect(documents[best_match]).to eq("washington")
-            end
+        it "finds the best match" do
+          VCR.use_cassette(cassette) do
+            expect(documents[best_match]).to eq("washington")
           end
         end
+      end
 
-        context "with engine: davinci" do
-          let(:engine) { "davinci" }
+      context "with engine: davinci" do
+        let(:engine) { "davinci" }
 
-          it "finds the best match" do
-            VCR.use_cassette(cassette) do
-              expect(documents[best_match]).to eq("washington")
-            end
+        it "finds the best match" do
+          VCR.use_cassette(cassette) do
+            expect(documents[best_match]).to eq("washington")
           end
         end
       end