Allow k to be higher than doc size in max_marginal_relevance_search (l…

…angchain-ai#1187)

Fixes issue langchain-ai#1186. For some reason, langchain-ai#1117 didn't seem to fix it.

langchain/vectorstores/faiss.py

@@ -188,10 +188,10 @@ def max_marginal_relevance_search_by_vector(
         selected_indices = [indices[0][i] for i in mmr_selected]
         docs = []
         for i in selected_indices:
-            _id = self.index_to_docstore_id[i]
-            if _id == -1:
+            if i == -1:
                 # This happens when not enough docs are returned.
                 continue
+            _id = self.index_to_docstore_id[i]
             doc = self.docstore.search(_id)
             if not isinstance(doc, Document):
                 raise ValueError(f"Could not find document for id {_id}, got {doc}")

tests/integration_tests/vectorstores/test_faiss.py

@@ -44,6 +44,10 @@ def test_faiss_vector_sim() -> None:
     output = docsearch.similarity_search_by_vector(query_vec, k=1)
     assert output == [Document(page_content="foo")]
 
+    # make sure we can have k > docstore size
+    output = docsearch.max_marginal_relevance_search_by_vector(query_vec, k=10)
+    assert len(output) == len(texts)
+
 
 def test_faiss_with_metadatas() -> None:
     """Test end to end construction and search."""