Add support for bulk embedding in Ollama

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## [Unreleased]
 - Improve the Langchain::Tool::Database tool
 - Allow explictly setting tool_choice on the Assistant instance
+- Add support for bulk embedding in Ollama
 
 ## [0.15.3] - 2024-08-27
 - Fix OpenAI#embed when text-embedding-ada-002 is used

lib/langchain/llm/ollama.rb

@@ -218,8 +218,8 @@ def embed(
       top_p: nil
     )
       parameters = {
-        prompt: text,
-        model: model
+        model: model,
+        input: Array(text)
       }.compact
 
       llm_parameters = {
@@ -243,7 +243,7 @@ def embed(
 
       parameters[:options] = llm_parameters.compact
 
-      response = client.post("api/embeddings") do |req|
+      response = client.post("api/embed") do |req|
         req.body = parameters
       end
 

lib/langchain/llm/response/ollama_response.rb

@@ -28,7 +28,7 @@ def embedding
     end
 
     def embeddings
-      [raw_response&.dig("embedding")]
+      raw_response&.dig("embeddings") || []
     end
 
     def role

spec/langchain/llm/ollama_spec.rb

@@ -25,7 +25,7 @@
 
   describe "#embed" do
     let(:response_body) {
-      {"embedding" => [0.1, 0.2, 0.3]}
+      {"embeddings" => [[0.1, 0.2, 0.3]]}
     }
 
     before do
@@ -36,6 +36,26 @@
       expect(subject.embed(text: "Hello, world!")).to be_a(Langchain::LLM::OllamaResponse)
       expect(subject.embed(text: "Hello, world!").embedding.count).to eq(3)
     end
+
+    it "sends an input array to the client" do
+      subject.embed(text: "Hello, world!")
+
+      expect(client).to have_received(:post) do |path, &block|
+        req = double("request").as_null_object
+        block.call(req)
+        expect(req).to have_received(:body=).with(hash_including(input: ["Hello, world!"]))
+      end
+    end
+
+    context "when the JSON response contains no embeddings" do
+      let(:response_body) {
+        {"embeddings" => []}
+      }
+
+      it "#embedding returns nil" do
+        expect(subject.embed(text: "Hello, world!").embedding).to be nil
+      end
+    end
   end
 
   describe "#complete" do