From ccef8c8a09e7ef40c81fd488afeef8e225b9fa03 Mon Sep 17 00:00:00 2001
From: Johannes Wesch <johannes.wesch@aleph-alpha.com>
Date: Tue, 31 Oct 2023 12:53:32 +0100
Subject: [PATCH] updated search to use index parameter

---
 .../connectors/document_index.py                | 17 +++++++++++++++--
 .../retrievers/document_index_retriever.py      |  2 +-
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/intelligence_layer/connectors/document_index.py b/src/intelligence_layer/connectors/document_index.py
index eb8a5b879..37b628bc2 100644
--- a/src/intelligence_layer/connectors/document_index.py
+++ b/src/intelligence_layer/connectors/document_index.py
@@ -93,15 +93,16 @@ def list_documents(self, namespace: str, collection: str) -> Any:
         response.raise_for_status()
         return response.json()
 
-    def search(
+    def index_search(
         self,
         namespace: str,
         collection: str,
+        index: str,
         query: str,
         max_results: int,
         min_score: float,
     ) -> Any:
-        url = f"{self._base_document_index_url}/collections/{namespace}/{collection}/search"
+        url = f"{self._base_document_index_url}/collections/{namespace}/{collection}/indexes/{index}/search"
         data = {
             "query": [{"modality": "text", "text": query}],
             "max_results": max_results,
@@ -111,3 +112,15 @@ def search(
         response = requests.post(url, data=json.dumps(data), headers=self.headers)
         response.raise_for_status()
         return response.json()
+    
+    def asymmetric_search(
+        self,
+        namespace: str,
+        collection: str,
+        query: str,
+        max_results: int,
+        min_score: float,
+    ) -> Any:
+        return self.index_search(namespace, collection, "asymmetric", query, max_results, min_score)
+    
+
diff --git a/src/intelligence_layer/connectors/retrievers/document_index_retriever.py b/src/intelligence_layer/connectors/retrievers/document_index_retriever.py
index 306856071..e6d436f66 100644
--- a/src/intelligence_layer/connectors/retrievers/document_index_retriever.py
+++ b/src/intelligence_layer/connectors/retrievers/document_index_retriever.py
@@ -43,7 +43,7 @@ def __init__(
         self._threshold = threshold
 
     def get_relevant_documents_with_scores(self, query: str) -> Sequence[SearchResult]:
-        response = self._document_index.search(
+        response = self._document_index.asymmetric_search(
             self._namespace, self._collection, query, self._k, self._threshold
         )
         relevant_chunks = [