diff --git a/libs/community/langchain_community/retrievers/google_vertex_ai_search.py b/libs/community/langchain_community/retrievers/google_vertex_ai_search.py index 78595a00cb..9544cc5817 100644 --- a/libs/community/langchain_community/retrievers/google_vertex_ai_search.py +++ b/libs/community/langchain_community/retrievers/google_vertex_ai_search.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple from langchain_core.callbacks import CallbackManagerForRetrieverRun from langchain_core.documents import Document @@ -345,6 +345,11 @@ class GoogleVertexAISearchRetriever(BaseRetriever, _BaseGoogleVertexAISearchRetr self, query: str, *, run_manager: CallbackManagerForRetrieverRun ) -> List[Document]: """Get documents relevant for a query.""" + return self.get_relevant_documents_with_response(query)[0] + + def get_relevant_documents_with_response( + self, query: str + ) -> Tuple[List[Document], Any]: from google.api_core.exceptions import InvalidArgument search_request = self._create_search_request(query) @@ -382,7 +387,7 @@ class GoogleVertexAISearchRetriever(BaseRetriever, _BaseGoogleVertexAISearchRetr + f" Got {self.engine_data_type}" ) - return documents + return documents, response class GoogleVertexAIMultiTurnSearchRetriever(