From edc471d1e2098bd9cdd62033d503431ca9d0a35e Mon Sep 17 00:00:00 2001
From: Diwank Singh Tomer <diwank.singh@gmail.com>
Date: Tue, 29 Oct 2024 00:44:23 +0530
Subject: [PATCH] Update search_docs_by_embedding.py

---
 agents-api/agents_api/models/docs/search_docs_by_embedding.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/agents-api/agents_api/models/docs/search_docs_by_embedding.py b/agents-api/agents_api/models/docs/search_docs_by_embedding.py
index 7f8bdaa61..382238f99 100644
--- a/agents-api/agents_api/models/docs/search_docs_by_embedding.py
+++ b/agents-api/agents_api/models/docs/search_docs_by_embedding.py
@@ -213,6 +213,8 @@ def search_docs_by_embedding(
                 distance = cos_dist(query, embedding),
                 distance <= {radius}
 
+            :limit {k*(3 if mmr_strength else 1)}   # Get more candidates for diversity
+
             :create _search_result {{
                 doc_id,
                 index,
@@ -244,7 +246,7 @@ def search_docs_by_embedding(
 
         # Sort the results by distance to find the closest matches
         :sort -mmr_score
-        :limit {k}
+        :limit {k*(3 if mmr_strength else 1)},   # Get more candidates for diversity
 
         :create _interim {{
             owner_type,