From 0874cba94d30dde34f4cc4866cc543dd85809e74 Mon Sep 17 00:00:00 2001
From: Andrew <morro141@gmail.com>
Date: Mon, 21 Aug 2023 18:19:05 -0400
Subject: [PATCH] Add k to pinecone's "ask" functionality (#297)

* update specs to allow for k value

* add k value to ask

* add k to similarity search
---
 lib/langchain/vectorsearch/pinecone.rb       | 4 ++--
 spec/langchain/vectorsearch/pinecone_spec.rb | 7 ++++---
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/lib/langchain/vectorsearch/pinecone.rb b/lib/langchain/vectorsearch/pinecone.rb
index b6d0bb313..864d2a538 100644
--- a/lib/langchain/vectorsearch/pinecone.rb
+++ b/lib/langchain/vectorsearch/pinecone.rb
@@ -141,8 +141,8 @@ def similarity_search_by_vector(embedding:, k: 4, namespace: "", filter: nil)
     # @param filter [String] The filter to use
     # @yield [String] Stream responses back one String at a time
     # @return [String] The answer to the question
-    def ask(question:, namespace: "", filter: nil, &block)
-      search_results = similarity_search(query: question, namespace: namespace, filter: filter)
+    def ask(question:, namespace: "", filter: nil, k: 4, &block)
+      search_results = similarity_search(query: question, namespace: namespace, filter: filter, k: k)
 
       context = search_results.map do |result|
         result.dig("metadata").to_s
diff --git a/spec/langchain/vectorsearch/pinecone_spec.rb b/spec/langchain/vectorsearch/pinecone_spec.rb
index 2bd8df051..631f084a2 100644
--- a/spec/langchain/vectorsearch/pinecone_spec.rb
+++ b/spec/langchain/vectorsearch/pinecone_spec.rb
@@ -305,11 +305,12 @@
     let(:question) { "How many times is \"lorem\" mentioned in this text?" }
     let(:prompt) { "Context:\n#{metadata}\n---\nQuestion: #{question}\n---\nAnswer:" }
     let(:answer) { "5 times" }
+    let(:k) { 4 }
 
     describe "without a namespace" do
       before do
         allow(subject).to receive(:similarity_search).with(
-          query: question, namespace: "", filter: nil
+          query: question, namespace: "", filter: nil, k: k
         ).and_return(matches)
         allow(subject.llm).to receive(:chat).with(prompt: prompt).and_return(answer)
       end
@@ -322,7 +323,7 @@
     describe "with a namespace" do
       before do
         allow(subject).to receive(:similarity_search).with(
-          query: question, namespace: namespace, filter: nil
+          query: question, namespace: namespace, filter: nil, k: k
         ).and_return(matches)
         allow(subject.llm).to receive(:chat).with(prompt: prompt).and_return(answer)
       end
@@ -335,7 +336,7 @@
     describe "with a filter" do
       before do
         allow(subject).to receive(:similarity_search).with(
-          query: question, namespace: "", filter: filter
+          query: question, namespace: "", filter: filter, k: k
         ).and_return(matches)
         allow(subject.llm).to receive(:chat).with(prompt: prompt).and_return(answer)
       end