From 5ab6b3909859998d31d601016dbeceb283de7e5b Mon Sep 17 00:00:00 2001 From: hsuyuming Date: Sun, 31 Mar 2024 14:54:56 -0600 Subject: [PATCH] community[patch]: add attribution_token within GoogleVertexAISearchRetriever (#18520) - **Description:** Add attribution_token within GoogleVertexAISearchRetriever so user can provide this information to Google support team or product team during debug session. Reference: https://cloud.google.com/generative-ai-app-builder/docs/view-analytics#user-events Attribution tokens. Attribution tokens are unique IDs generated by Vertex AI Search and returned with each search request. Make sure to include that attribution token as UserEvent.attributionToken with any user events resulting from a search. This is needed to identify if a search is served by the API. Only user events with a Google-generated attribution token are used to compute metrics. - **Issue:** No - **Dependencies:** No - **Twitter handle:** abehsu1992626 --------- Co-authored-by: Bagatur --- .../retrievers/google_vertex_ai_search.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) 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(