From c96116bd1435ea1c6a942eaf4a70b7beaf007b5c Mon Sep 17 00:00:00 2001
From: "ravindu.somawansa" <ravindu.somawansa@airliquide.com>
Date: Fri, 12 Jul 2024 16:23:28 +0200
Subject: [PATCH] fix: removed argument nextToken

---
 libs/aws/langchain_aws/retrievers/bedrock.py | 36 +-------------------
 1 file changed, 1 insertion(+), 35 deletions(-)

diff --git a/libs/aws/langchain_aws/retrievers/bedrock.py b/libs/aws/langchain_aws/retrievers/bedrock.py
index 293c0b1f..bab024f9 100644
--- a/libs/aws/langchain_aws/retrievers/bedrock.py
+++ b/libs/aws/langchain_aws/retrievers/bedrock.py
@@ -160,12 +160,11 @@ def _get_relevant_documents(
         query: str,
         *,
         run_manager: CallbackManagerForRetrieverRun,
-        next_token: Optional[str] = None,
     ) -> List[Document]:
         response = self.client.retrieve(
             retrievalQuery={"text": query.strip()},
             knowledgeBaseId=self.knowledge_base_id,
-            nextToken=next_token,
+            nextToken=None,
             retrievalConfiguration=self.retrieval_config.dict(),
         )
         results = response["retrievalResults"]
@@ -185,36 +184,3 @@ def _get_relevant_documents(
             )
 
         return self._filter_by_score_confidence(docs=documents)
-
-    def _get_relevant_documents_with_token(
-        self,
-        query: str,
-        *,
-        run_manager: CallbackManagerForRetrieverRun,
-        next_token: Optional[str] = None,
-    ) -> Tuple[List[Document], Optional[str]]:
-        response = self.client.retrieve(
-            retrievalQuery={"text": query.strip()},
-            knowledgeBaseId=self.knowledge_base_id,
-            nextToken=next_token or self.retrieval_config.nextToken,
-            retrievalConfiguration=self.retrieval_config.dict(),
-        )
-        new_next_token = response.get("nextToken", None)
-        results = response["retrievalResults"]
-        documents = []
-        for result in results:
-            content = result["content"]["text"]
-            result.pop("content")
-            if "score" not in result:
-                result["score"] = 0
-            if "metadata" in result:
-                result["source_metadata"] = result.pop("metadata")
-            documents.append(
-                Document(
-                    page_content=content,
-                    metadata=result,
-                )
-            )
-
-        filtered_documents = self._filter_by_score_confidence(docs=documents)
-        return (filtered_documents, new_next_token)