Skip to content

Commit

Permalink
fix: removed argument nextToken
Browse files Browse the repository at this point in the history
  • Loading branch information
ravindu.somawansa committed Jul 12, 2024
1 parent 0e7ede7 commit c96116b
Showing 1 changed file with 1 addition and 35 deletions.
36 changes: 1 addition & 35 deletions libs/aws/langchain_aws/retrievers/bedrock.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand All @@ -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)

0 comments on commit c96116b

Please sign in to comment.