From 7decd8ab31be319e71747546dd85572ec8b3772f Mon Sep 17 00:00:00 2001
From: Diwank Singh Tomer <diwank.singh@gmail.com>
Date: Tue, 15 Oct 2024 14:03:57 -0400
Subject: [PATCH] fix(agents-api): Fix search query to restrict to
 candidate[doc_ids]

Signed-off-by: Diwank Singh Tomer <diwank.singh@gmail.com>
---
 .../agents_api/models/docs/search_docs_by_embedding.py      | 6 ++++++
 agents-api/agents_api/models/docs/search_docs_by_text.py    | 1 +
 2 files changed, 7 insertions(+)

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 856900709..90db8ac26 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
@@ -99,6 +99,7 @@ def search_docs_by_embedding(
             index1,
             min(dist)
         ] :=
+            candidate[doc_id],
             *snippets {{
                 doc_id,
                 index: index1,
@@ -115,15 +116,18 @@ def search_docs_by_embedding(
             dist = cos_dist(embedding1, embedding2)
 
         doclength[doc_id, max(index)] :=
+            candidate[doc_id],
             *snippets {{
                 doc_id,
                 index,
             }}
 
         get_intersnippet[doc_id, index, distance] :=
+            candidate[doc_id],
             intersnippet_distance[doc_id, _, distance]
 
         get_intersnippet[doc_id, index, distance] :=
+            candidate[doc_id],
             not intersnippet_distance[doc_id, _, distance],
             distance = 0.0
 
@@ -153,6 +157,7 @@ def search_docs_by_embedding(
             distance,
             mmr_score,
         ] :=
+            candidate[doc_id],
             search_result[doc_id, content, index, distance],
             get_intersnippet[doc_id, index, intersnippet_distance],
             mmr_score = {mmr_lambda} * (distance - (1.0 - {mmr_lambda}) * intersnippet_distance),
@@ -167,6 +172,7 @@ def search_docs_by_embedding(
             mmr_score,
             title,
         ] := 
+            candidate[doc_id],
             *docs {{
                 owner_type,
                 owner_id,
diff --git a/agents-api/agents_api/models/docs/search_docs_by_text.py b/agents-api/agents_api/models/docs/search_docs_by_text.py
index bb700a494..bc7063e87 100644
--- a/agents-api/agents_api/models/docs/search_docs_by_text.py
+++ b/agents-api/agents_api/models/docs/search_docs_by_text.py
@@ -154,6 +154,7 @@ def search_docs_by_text(
             distance,
             title,
         ] := 
+            candidate[id],
             input[owner_type, owner_id],
             m[
                 id,