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 <baskaryan@gmail.com>
pull/19649/head
hsuyuming 3 months ago committed by GitHub
parent f98d7f7494
commit 5ab6b39098
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -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(

Loading…
Cancel
Save