diff --git a/libs/langchain/langchain/retrievers/google_cloud_enterprise_search.py b/libs/langchain/langchain/retrievers/google_cloud_enterprise_search.py index 2acb01b363..b609317a33 100644 --- a/libs/langchain/langchain/retrievers/google_cloud_enterprise_search.py +++ b/libs/langchain/langchain/retrievers/google_cloud_enterprise_search.py @@ -106,34 +106,23 @@ class GoogleCloudEnterpriseSearchRetriever(BaseRetriever): self, results: Sequence[SearchResult] ) -> List[Document]: """Converts a sequence of search results to a list of LangChain documents.""" - from google.protobuf.json_format import MessageToDict + documents: List[Document] = [] - documents = [] for result in results: - document_dict = MessageToDict(result.document._pb) - derived_struct_data = document_dict.get("derivedStructData", None) - if derived_struct_data: - doc_metadata = document_dict.get("structData", {}) - chunk_type = ( - "extractive_answers" - if self.get_extractive_answers - else "extractive_segments" + derived_struct_data = result.document.derived_struct_data + doc_metadata = result.document.struct_data + doc_metadata.source = derived_struct_data.link or "" + doc_metadata.id = result.document.id + + for chunk in ( + derived_struct_data.extractive_answers + or derived_struct_data.extractive_segments + ): + if hasattr(chunk, "page_number"): + doc_metadata.source += f":{chunk.page_number}" + documents.append( + Document(page_content=chunk.content, metadata=doc_metadata) ) - for chunk in derived_struct_data.get(chunk_type, []): - if chunk_type == "extractive_answers": - doc_metadata["source"] = ( - f"{derived_struct_data.get('link', '')}" - f":{chunk.get('pageNumber', '')}" - ) - else: - doc_metadata[ - "source" - ] = f"{derived_struct_data.get('link', '')}" - doc_metadata["id"] = document_dict["id"] - document = Document( - page_content=chunk.get("content", ""), metadata=doc_metadata - ) - documents.append(document) return documents @@ -162,7 +151,7 @@ class GoogleCloudEnterpriseSearchRetriever(BaseRetriever): extractive_content_spec=extractive_content_spec, ) - request = SearchRequest( + return SearchRequest( query=query, filter=self.filter, serving_config=self._serving_config, @@ -171,8 +160,6 @@ class GoogleCloudEnterpriseSearchRetriever(BaseRetriever): query_expansion_spec=query_expansion_spec, ) - return request - def _get_relevant_documents( self, query: str, *, run_manager: CallbackManagerForRetrieverRun ) -> List[Document]: