From a7af24ad694c3b61c4bb1d46e7037e99410530f2 Mon Sep 17 00:00:00 2001 From: sailobo Date: Tue, 8 Sep 2026 13:19:42 -0600 Subject: [PATCH 1/4] Add files via upload Adding code samples for Agent Search --- .../converse_conversation_sample.py | 126 ++++ .../create_document_metadata_sample.py | 111 ++++ .../enable_gemini_layout_parser_sample.py | 114 ++++ .../import_custom_chunks_sample.py | 91 +++ discoveryengine/poll_lro_robust_sample.py | 118 ++++ discoveryengine/search_chunks_sample.py | 131 +++++ .../search_extractive_segments_sample.py | 141 +++++ .../supportability_samples_test.py | 545 ++++++++++++++++++ .../update_data_store_gemini_parser_sample.py | 97 ++++ discoveryengine/write_user_event_sample.py | 129 +++++ 10 files changed, 1603 insertions(+) create mode 100644 discoveryengine/converse_conversation_sample.py create mode 100644 discoveryengine/create_document_metadata_sample.py create mode 100644 discoveryengine/enable_gemini_layout_parser_sample.py create mode 100644 discoveryengine/import_custom_chunks_sample.py create mode 100644 discoveryengine/poll_lro_robust_sample.py create mode 100644 discoveryengine/search_chunks_sample.py create mode 100644 discoveryengine/search_extractive_segments_sample.py create mode 100644 discoveryengine/supportability_samples_test.py create mode 100644 discoveryengine/update_data_store_gemini_parser_sample.py create mode 100644 discoveryengine/write_user_event_sample.py diff --git a/discoveryengine/converse_conversation_sample.py b/discoveryengine/converse_conversation_sample.py new file mode 100644 index 00000000000..4bdffe01258 --- /dev/null +++ b/discoveryengine/converse_conversation_sample.py @@ -0,0 +1,126 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Multi-turn conversational search sample for Agent Search.""" + +# [START genappbuilder_converse_conversation] +from typing import Optional + +from google.cloud import discoveryengine_v1 as discoveryengine + + +def multi_turn_conversational_search( + project_id: str, + location: str, + data_store_id: str, + query_text: str, + conversation_id: Optional[str] = None, + user_pseudo_id: str = "user_pseudo_id_12345", +) -> discoveryengine.ConverseConversationResponse: + """Performs a multi-turn conversational search with session management and citation parsing. + + Multi-turn conversational search maintains conversational context across + user interactions: + 1. Create a Conversation session with + `ConversationalSearchServiceClient.create_conversation()`. + 2. Pass the conversation `name` in subsequent `ConverseConversationRequest` + calls. + 3. Parse citations and search results to display grounded source references. + + Args: + project_id: Google Cloud project ID or project number. + location: Data store location (e.g., 'global', 'us', 'eu'). + data_store_id: Data store ID. + query_text: Natural language user question or follow-up query. + conversation_id: Optional existing conversation ID to continue a session. + user_pseudo_id: Unique visitor/user identifier. + + Returns: + ConverseConversationResponse containing the answer and grounded citations. + """ + client = discoveryengine.ConversationalSearchServiceClient() + + parent = ( + f"projects/{project_id}/locations/{location}/collections/default_collection" + f"/dataStores/{data_store_id}" + ) + + serving_config = f"{parent}/servingConfigs/default_search" + + # Step 1: Create a new Conversation session if not provided + if not conversation_id: + conversation = discoveryengine.Conversation( + user_pseudo_id=user_pseudo_id, + state=discoveryengine.Conversation.State.IN_PROGRESS, + ) + created_conversation = client.create_conversation( + parent=parent, + conversation=conversation, + ) + conversation_name = created_conversation.name + print(f"Created new conversation session: {conversation_name}") + else: + conversation_name = ( + f"{parent}/conversations/{conversation_id}" + if not conversation_id.startswith("projects/") + else conversation_id + ) + print(f"Continuing existing conversation session: {conversation_name}") + + # Step 2: Send query in the conversation + query_input = discoveryengine.TextInput(input=query_text) + + summary_spec = discoveryengine.SearchRequest.ContentSearchSpec.SummarySpec( + include_citations=True + ) + + request = discoveryengine.ConverseConversationRequest( + name=conversation_name, + query=query_input, + serving_config=serving_config, + summary_spec=summary_spec, + ) + + response = client.converse_conversation(request=request) + + reply_text = "" + if response.reply: + if ( + hasattr(response.reply, "summary") + and response.reply.summary + and response.reply.summary.summary_text + ): + reply_text = response.reply.summary.summary_text + elif hasattr(response.reply, "reply") and response.reply.reply: + reply_text = response.reply.reply + + print(f"\nUser Query: {query_text}") + print(f"AI Generated Reply: {reply_text}") + + # Step 3: Parse and display source citations and grounding results + print("\nGrounding Citations & Sources:") + for idx, search_result in enumerate(response.search_results, 1): + doc = search_result.document + print(f" [{idx}] Document ID: {doc.id}") + struct_data = doc.derived_struct_data or doc.struct_data + if struct_data: + title = struct_data.get("title", "No Title") + link = struct_data.get("link", "No Link") + print(f" Title: {title}") + print(f" URI: {link}") + + return response + + +# [END genappbuilder_converse_conversation] diff --git a/discoveryengine/create_document_metadata_sample.py b/discoveryengine/create_document_metadata_sample.py new file mode 100644 index 00000000000..ec2fb1cb617 --- /dev/null +++ b/discoveryengine/create_document_metadata_sample.py @@ -0,0 +1,111 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Document metadata ingestion sample for Agent Search.""" + +# [START genappbuilder_create_document_metadata] +from typing import List, Optional + +from google.cloud import discoveryengine_v1 as discoveryengine + + +def create_structured_document_with_metadata( + project_id: str, + location: str, + data_store_id: str, + document_id: str, + title: str, + uri: str, + category: str, + rating: float, + tags: List[str], + content_text: Optional[str] = None, +) -> discoveryengine.Document: + """Creates a document with custom structured metadata and content. + + When ingesting documents into Agent Search, passing authoritative + metadata in `struct_data` provides key benefits: + 1. Prevents URI corruption: LLM layout parsers often strip underscores + from URLs during markdown translation. Explicitly passing `uri` in + `struct_data` guarantees exact URL preservation for frontend rendering. + 2. Schema Alignment: Explicit metadata attributes (categories, numeric + ratings, tags) enable exact filtering and faceting in search queries. + 3. Chunking Inheritance: When ingested into a data store with document + chunking enabled, `struct_data` is preserved on the parent document. In + chunk search results (`searchResultMode=CHUNKS`), access metadata via + `result.chunk.document_metadata.struct_data`. + + Args: + project_id: Google Cloud project ID or project number. + location: Data store location (e.g., 'global', 'us', 'eu'). + data_store_id: Target data store ID. + document_id: Unique document identifier. + title: Title of the document. + uri: Original document URL or Cloud Storage URI (exact string preserved). + category: Document taxonomy/category for faceted filtering. + rating: Numerical score/rating for numerical filtering. + tags: List of string tags/keywords. + content_text: Optional raw text body for full-text indexing. + + Returns: + The created Document proto. + """ + client = discoveryengine.DocumentServiceClient() + + # Document branch 0 is the default serving branch + parent = ( + f"projects/{project_id}/locations/{location}/collections/default_collection" + f"/dataStores/{data_store_id}/branches/0" + ) + + # Build structured metadata dictionary + metadata = { + "title": title, + "url": uri, + "category": category, + "rating": rating, + "tags": tags, + } + + document = discoveryengine.Document( + id=document_id, + struct_data=metadata, + ) + + # Optional unstructured text content + if content_text: + document.content = discoveryengine.Document.Content( + mime_type="text/plain", + raw_bytes=content_text.encode("utf-8"), + ) + + request = discoveryengine.CreateDocumentRequest( + parent=parent, + document=document, + document_id=document_id, + ) + + response = client.create_document(request=request) + + print(f"Created Document ID: {response.id}") + print(f" Name: {response.name}") + print(f" Metadata URL (exact): {response.struct_data.get('url')}") + print(f" Category: {response.struct_data.get('category')}") + print(f" Rating: {response.struct_data.get('rating')}") + print(f" Tags: {response.struct_data.get('tags')}") + + return response + + +# [END genappbuilder_create_document_metadata] diff --git a/discoveryengine/enable_gemini_layout_parser_sample.py b/discoveryengine/enable_gemini_layout_parser_sample.py new file mode 100644 index 00000000000..132250a97b4 --- /dev/null +++ b/discoveryengine/enable_gemini_layout_parser_sample.py @@ -0,0 +1,114 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Gemini Advanced Layout Parser creation sample for Agent Search.""" + +# [START genappbuilder_enable_gemini_layout_parser] +from google.api_core.client_options import ClientOptions +from google.cloud import discoveryengine_v1beta as discoveryengine + + +def create_data_store_with_gemini_parser( + project_id: str, + location: str, + data_store_id: str, + display_name: str, + enable_table_annotation: bool = True, + enable_image_annotation: bool = True, +) -> discoveryengine.DataStore: + """Creates a DataStore configured with Gemini Advanced Layout Parser (Pre-GA / Preview). + + Gemini layout parsing ('enable_llm_layout_parsing=True') uses Gemini + multimodal + models to provide superior table extraction, reading order analysis, and + optical + character recognition on PDFs. When combined with 'LayoutBasedChunkingConfig', + it ensures structural elements like tables, lists, and sections are parsed + cleanly + for downstream RAG and answer generation. + + Args: + project_id: Google Cloud project ID. + location: DataStore location (e.g., 'global', 'us', 'eu'). + data_store_id: Unique identifier for the DataStore. + display_name: Human-readable display name for the DataStore. + enable_table_annotation: Whether to generate LLM descriptions for + extracted tables. + enable_image_annotation: Whether to generate LLM descriptions for + extracted images. + + Returns: + The created DataStore object (or the Long-Running Operation result). + """ + client_options = ( + ClientOptions(api_endpoint=f"{location}-discoveryengine.googleapis.com") + if location != "global" + else None + ) + client = discoveryengine.DataStoreServiceClient(client_options=client_options) + + parent = f"projects/{project_id}/locations/{location}/collections/default_collection" + + # 1. Configure Layout Parsing with Gemini LLM Enhancement + layout_parsing_config = discoveryengine.DocumentProcessingConfig.ParsingConfig.LayoutParsingConfig( + enable_llm_layout_parsing=True, # Enables Gemini LLM-based layout parsing + enable_table_annotation=enable_table_annotation, + enable_image_annotation=enable_image_annotation, + ) + + parsing_config = discoveryengine.DocumentProcessingConfig.ParsingConfig( + layout_parsing_config=layout_parsing_config + ) + + # 2. Configure Layout-Based Chunking for RAG + chunking_config = discoveryengine.DocumentProcessingConfig.ChunkingConfig( + layout_based_chunking_config=discoveryengine.DocumentProcessingConfig.ChunkingConfig.LayoutBasedChunkingConfig( + chunk_size=500, + include_ancestor_headings=True, + ) + ) + + # 3. Assemble DocumentProcessingConfig + doc_processing_config = discoveryengine.DocumentProcessingConfig( + default_parsing_config=parsing_config, + chunking_config=chunking_config, + ) + + # 4. Construct DataStore + data_store = discoveryengine.DataStore( + display_name=display_name, + industry_vertical=discoveryengine.IndustryVertical.GENERIC, + solution_types=[discoveryengine.SolutionType.SOLUTION_TYPE_SEARCH], + content_config=discoveryengine.DataStore.ContentConfig.CONTENT_REQUIRED, + document_processing_config=doc_processing_config, + ) + + request = discoveryengine.CreateDataStoreRequest( + parent=parent, + data_store=data_store, + data_store_id=data_store_id, + ) + + operation = client.create_data_store(request=request) + print(f"Waiting for DataStore creation operation: {operation.operation.name}") + created_data_store = operation.result() + + print("Successfully created DataStore with Gemini Layout Parser:") + print(f" Name: {created_data_store.name}") + print(f" Display Name: {created_data_store.display_name}") + + return created_data_store + + +# [END genappbuilder_enable_gemini_layout_parser] diff --git a/discoveryengine/import_custom_chunks_sample.py b/discoveryengine/import_custom_chunks_sample.py new file mode 100644 index 00000000000..8776614ca25 --- /dev/null +++ b/discoveryengine/import_custom_chunks_sample.py @@ -0,0 +1,91 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Bring Your Own Chunks (BYOC) import sample for Agent Search.""" + +# [START genappbuilder_import_custom_chunks] +from google.api_core.operation import Operation +from google.cloud import discoveryengine_v1 as discoveryengine + + +def import_custom_chunk_documents( + project_id: str, + location: str, + data_store_id: str, + gcs_uri: str, +) -> Operation: + """Imports pre-chunked JSON documents from Cloud Storage into an Agent Search data store. + + When bringing pre-chunked documents (e.g., from LangChain, LlamaIndex, or + custom semantic splitters), the JSON files in Cloud Storage must adhere to the + Bring Your Own Chunks (BYOC) schema with a top-level `documentMetadata` object + and a `chunks` array: + + ```json + { + "documentMetadata": { + "title": "Example Document Title", + "uri": "https://example.com/doc.pdf", + "structData": { + "category": "technical", + "rating": 4.5 + } + }, + "chunks": [ + { + "id": "chunk_0", + "content": "First section text...", + "pageSpan": { + "pageStart": 1, + "pageEnd": 1 + } + } + ] + } + ``` + + Args: + project_id: Google Cloud project ID or project number. + location: Data store location (e.g., 'global', 'us', 'eu'). + data_store_id: Target data store ID. + gcs_uri: Cloud Storage URI of BYOC JSON files (e.g., + 'gs://my-bucket/chunks/*.json'). + + Returns: + The Long-Running Operation (Operation) for the document import. + """ + client = discoveryengine.DocumentServiceClient() + + # Document branch 0 is the default serving branch + parent = ( + f"projects/{project_id}/locations/{location}/collections/default_collection" + f"/dataStores/{data_store_id}/branches/0" + ) + + request = discoveryengine.ImportDocumentsRequest( + parent=parent, + gcs_source=discoveryengine.GcsSource( + input_uris=[gcs_uri], + data_schema="custom", + ), + # FULL reconciliation replaces the dataset; INCREMENTAL adds/updates + reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL, + ) + + operation = client.import_documents(request=request) + print(f"Triggered BYOC Document Import LRO: {operation.operation.name}") + return operation + + +# [END genappbuilder_import_custom_chunks] diff --git a/discoveryengine/poll_lro_robust_sample.py b/discoveryengine/poll_lro_robust_sample.py new file mode 100644 index 00000000000..2a2ca5edecb --- /dev/null +++ b/discoveryengine/poll_lro_robust_sample.py @@ -0,0 +1,118 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Robust Long-Running Operation (LRO) polling sample for Agent Search.""" + +# [START genappbuilder_poll_lro_robust] +import random +import time + +from google.cloud import discoveryengine_v1 as discoveryengine +from google.longrunning import operations_pb2 + + +def poll_long_running_operation_robust( + operation_name: str, + initial_delay_seconds: float = 5.0, + max_delay_seconds: float = 60.0, + backoff_factor: float = 1.5, + timeout_seconds: float = 600.0, +) -> operations_pb2.Operation: + """Polls a long-running operation with exponential backoff and error diagnosis. + + Asynchronous operations (e.g., document ingestion, schema updates, data store + provisioning) often take several minutes. Best practices for reliable polling + include: + 1. Exponential backoff and jitter to prevent quota exhaustion and rate + limiting. + 2. Detailed error inspection when `operation.done` is True and + `operation.error` + contains failure status details. + 3. Inspecting metadata for partial failure counts and GCS error logs. + + Args: + operation_name: Full resource name of the operation (e.g., + 'projects/.../locations/.../operations/...'). + initial_delay_seconds: Starting backoff delay in seconds. + max_delay_seconds: Maximum backoff delay cap. + backoff_factor: Multiplier for exponential backoff. + timeout_seconds: Maximum total duration before raising a TimeoutError. + + Returns: + The completed operations_pb2.Operation proto. + + Raises: + TimeoutError: If the operation does not finish within `timeout_seconds`. + RuntimeError: If the operation finishes with an error status. + """ + # Use DocumentServiceClient's underlying operations client or OperationsClient + client = discoveryengine.DocumentServiceClient() + operations_client = client.transport.operations_client + + start_time = time.time() + current_delay = initial_delay_seconds + attempt = 1 + + print(f"Starting robust polling for operation: {operation_name}") + + while True: + elapsed = time.time() - start_time + if elapsed > timeout_seconds: + raise TimeoutError( + f"Operation '{operation_name}' exceeded timeout of {timeout_seconds}" + " seconds." + ) + + operation = operations_client.get_operation(name=operation_name) + + if operation.done: + print(f"\n[Attempt {attempt}] Operation completed in {elapsed:.1f}s!") + + # 1. Check for fatal operation error + if operation.HasField("error") and ( + operation.error.code != 0 or operation.error.message + ): + error = operation.error + raise RuntimeError( + f"Operation failed with Code {error.code}: {error.message}" + ) + + # 2. Check for operation metadata details + if operation.HasField("metadata"): + print("Operation Metadata Details:") + print(f" Type: {operation.metadata.type_url}") + + # 3. Check for operation response + if operation.HasField("response"): + print("Operation Response:") + print(f" Type: {operation.response.type_url}") + + return operation + + print( + f"[Attempt {attempt} - {elapsed:.1f}s elapsed] Operation still in" + f" progress... Waiting {current_delay:.1f}s." + ) + + time.sleep(current_delay) + + # Exponential backoff with jitter (+/- 10%) + jitter = random.uniform(0.9, 1.1) + current_delay = min( + current_delay * backoff_factor * jitter, max_delay_seconds + ) + attempt += 1 + + +# [END genappbuilder_poll_lro_robust] diff --git a/discoveryengine/search_chunks_sample.py b/discoveryengine/search_chunks_sample.py new file mode 100644 index 00000000000..2bd624d0319 --- /dev/null +++ b/discoveryengine/search_chunks_sample.py @@ -0,0 +1,131 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Search chunks and parent metadata sample for Agent Search.""" + +# [START genappbuilder_search_chunks] +from typing import Optional + +from google.cloud import discoveryengine_v1 as discoveryengine + + +def search_chunks_with_metadata( + project_id: str, + location: str, + data_store_id: str, + search_query: str, + num_previous_chunks: int = 1, + num_next_chunks: int = 1, + filter_expr: Optional[str] = None, + page_size: int = 5, +) -> discoveryengine.SearchResponse: + """Searches a chunk-enabled data store and parses chunk content and parent metadata. + + When searching a data store with document chunking enabled: + 1. Search results populate `result.chunk` rather than `result.document`. + 2. Parent document metadata (such as `title`, preserved `uri`, or custom + `struct_data`) must be retrieved via `result.chunk.document_metadata`. + 3. `chunk_spec` enables windowing with adjacent chunks (`previous_chunks`, + `next_chunks`) to provide broader context to downstream LLM prompts without + losing fine-grained relevance. + + Args: + project_id: Google Cloud project ID or project number. + location: Data store location (e.g., 'global', 'us', 'eu'). + data_store_id: Target data store ID. + search_query: The search query string. + num_previous_chunks: Number of preceding adjacent chunks to return (0-5). + num_next_chunks: Number of succeeding adjacent chunks to return (0-5). + filter_expr: Optional filter expression targeting parent structured + metadata. + page_size: Number of chunk results to return per page. + + Returns: + The SearchResponse proto. + """ + client = discoveryengine.SearchServiceClient() + + serving_config = ( + f"projects/{project_id}/locations/{location}/collections/default_collection" + f"/dataStores/{data_store_id}/servingConfigs/default_search" + ) + + # Configure search request for chunk retrieval with adjacent chunk context + content_search_spec = discoveryengine.SearchRequest.ContentSearchSpec( + search_result_mode=discoveryengine.SearchRequest.ContentSearchSpec.SearchResultMode.CHUNKS, + chunk_spec=discoveryengine.SearchRequest.ContentSearchSpec.ChunkSpec( + num_previous_chunks=num_previous_chunks, + num_next_chunks=num_next_chunks, + ), + ) + + request = discoveryengine.SearchRequest( + serving_config=serving_config, + query=search_query, + page_size=page_size, + content_search_spec=content_search_spec, + filter=filter_expr, + ) + + response = client.search(request=request) + + print(f"Search query: '{search_query}'") + print(f"Total results: {len(response.results)}") + print(f"Attribution Token: {response.attribution_token}") + + for idx, result in enumerate(response.results, start=1): + chunk = result.chunk + doc_metadata = chunk.document_metadata if chunk else None + + print(f"\n--- Result #{idx} ---") + if not chunk: + print(" (Result did not contain chunk data)") + continue + + print(f" Chunk ID: {chunk.id}") + print(f" Relevance Score: {chunk.relevance_score:.4f}") + + # Parent document metadata + if doc_metadata: + print(f" Parent Title: {doc_metadata.title}") + print(f" Parent URI: {doc_metadata.uri}") + if doc_metadata.struct_data: + print(f" Parent Structured Data: {dict(doc_metadata.struct_data)}") + + # Page span + if chunk.page_span: + print( + f" Page Span: {chunk.page_span.page_start} -" + f" {chunk.page_span.page_end}" + ) + + # Primary chunk text content + content_preview = ( + chunk.content[:150].replace("\n", " ") if chunk.content else "" + ) + print(f" Chunk Content: {content_preview}...") + + # Adjacent chunks for expanded LLM context + if chunk.chunk_metadata: + prev_count = len(chunk.chunk_metadata.previous_chunks) + next_count = len(chunk.chunk_metadata.next_chunks) + print( + f" Adjacent Context: {prev_count} previous chunk(s), {next_count}" + " next chunk(s)" + ) + + return response + + +# [END genappbuilder_search_chunks] diff --git a/discoveryengine/search_extractive_segments_sample.py b/discoveryengine/search_extractive_segments_sample.py new file mode 100644 index 00000000000..2e76d908d7e --- /dev/null +++ b/discoveryengine/search_extractive_segments_sample.py @@ -0,0 +1,141 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Search with extractive segments sample for Agent Search.""" + +# [START genappbuilder_search_extractive_segments] +from typing import List, Optional + +from google.cloud import discoveryengine_v1 as discoveryengine + + +def search_with_extractive_segments( + project_id: str, + location: str, + engine_id: str, + search_query: str, + data_store_ids: Optional[List[str]] = None, +) -> discoveryengine.SearchResponse: + """Performs a search query using extractive segments across one or more data stores. + + Using Extractive Segments (`max_extractive_segment_count`) enables + multi-datastore + blended search across chunked and unchunked data stores uniformly. + Extractive Segments work across all data store configurations and are + array-order + independent, returning exact source passages from indexed documents. + + Args: + project_id: Google Cloud project ID or project number. + location: Data store location (e.g., 'global', 'us', 'eu'). + engine_id: Search engine (app) ID. + search_query: The search text query. + data_store_ids: Optional list of specific data store IDs to blend. + + Returns: + SearchResponse containing results with extractive segments. + """ + client = discoveryengine.SearchServiceClient() + + serving_config = ( + f"projects/{project_id}/locations/{location}/collections/default_collection" + f"/engines/{engine_id}/servingConfigs/default_search" + ) + + # Configure Extractive Content Spec for segments + content_search_spec = discoveryengine.SearchRequest.ContentSearchSpec( + extractive_content_spec=discoveryengine.SearchRequest.ContentSearchSpec.ExtractiveContentSpec( + max_extractive_segment_count=3, + return_extractive_segment_score=True, + num_previous_segments=1, + num_next_segments=1, + ), + snippet_spec=discoveryengine.SearchRequest.ContentSearchSpec.SnippetSpec( + return_snippet=True + ), + ) + + # Optional multi-datastore scoping + data_store_specs = [] + if data_store_ids: + for ds_id in data_store_ids: + ds_path = ( + f"projects/{project_id}/locations/{location}" + f"/collections/default_collection/dataStores/{ds_id}" + ) + data_store_specs.append( + discoveryengine.SearchRequest.DataStoreSpec(data_store=ds_path) + ) + + request = discoveryengine.SearchRequest( + serving_config=serving_config, + query=search_query, + page_size=10, + content_search_spec=content_search_spec, + data_store_specs=data_store_specs if data_store_specs else None, + ) + + response = client.search(request=request) + + print(f"Search Results for query '{search_query}':") + print(f"Attribution Token: {response.attribution_token}") + print(f"Total Results: {response.total_size}") + + for result in response.results: + doc = result.document + print(f"\n- Document ID: {doc.id}") + print(f" Name: {doc.name}") + + derived_data = doc.derived_struct_data + raw_data = doc.struct_data + + # Retrieve title and URI from derived metadata or raw document struct_data + title = "No Title" + link = "No Link" + if derived_data: + title = derived_data.get("title") or title + link = derived_data.get("link") or derived_data.get("url") or link + if raw_data: + if title == "No Title": + title = raw_data.get("title") or title + if link == "No Link": + link = raw_data.get("link") or raw_data.get("url") or link + + print(f" Title: {title}") + print(f" Link: {link}") + + # Extract Extractive Segments + if derived_data and "extractive_segments" in derived_data: + segments = derived_data.get("extractive_segments", []) + for idx, segment in enumerate(segments, 1): + page_number = segment.get("pageNumber", "N/A") + relevance_score = segment.get("relevanceScore", "N/A") + content = segment.get("content", "") + print( + f" [Segment {idx}] (Page: {page_number}, Score:" + f" {relevance_score}):" + ) + print(f" {content.strip()}") + + # Extract HTML snippets with highlighting + if derived_data and "snippets" in derived_data: + snippets = derived_data.get("snippets", []) + for idx, snippet_entry in enumerate(snippets, 1): + snippet_text = snippet_entry.get("snippet", "") + print(f" [Snippet {idx}]: {snippet_text.strip()}") + + return response + + +# [END genappbuilder_search_extractive_segments] diff --git a/discoveryengine/supportability_samples_test.py b/discoveryengine/supportability_samples_test.py new file mode 100644 index 00000000000..546ef298396 --- /dev/null +++ b/discoveryengine/supportability_samples_test.py @@ -0,0 +1,545 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for Vertex AI Search / Agent Search supportability code samples.""" + +import os +import unittest +from unittest import mock + +from google.cloud import discoveryengine_v1 as discoveryengine +from google.longrunning import operations_pb2 +from google.protobuf import any_pb2 +from google.protobuf import struct_pb2 +from google.rpc import status_pb2 + +import converse_conversation_sample +import create_document_metadata_sample +import enable_gemini_layout_parser_sample +import import_custom_chunks_sample +import poll_lro_robust_sample +import search_chunks_sample +import search_extractive_segments_sample +import update_data_store_gemini_parser_sample +import write_user_event_sample + + +class SupportabilitySamplesTest(unittest.TestCase): + + @mock.patch("google.cloud.discoveryengine_v1.SearchServiceClient") + def test_search_with_extractive_segments(self, mock_client_cls): + mock_client = mock.MagicMock() + mock_client_cls.return_value = mock_client + + mock_doc = discoveryengine.Document( + id="doc-123", + name="projects/p/locations/l/collections/c/dataStores/ds/branches/0/documents/doc-123", + derived_struct_data={ + "title": "Alphabet 2024 Q4 Report", + "link": "https://example.com/reports/2024_q4.pdf", + "extractive_segments": [{ + "pageNumber": "12", + "relevanceScore": 0.94, + "content": "Cloud revenue increased 29% year over year.", + }], + "snippets": [{"snippet": "Cloud revenue increased 29%."}], + }, + ) + + mock_search_result = discoveryengine.SearchResponse.SearchResult( + id="doc-123", + document=mock_doc, + ) + + mock_response = discoveryengine.SearchResponse( + results=[mock_search_result], + attribution_token="test_attribution_token_123", + total_size=1, + ) + mock_client.search.return_value = mock_response + + response = ( + search_extractive_segments_sample.search_with_extractive_segments( + project_id="test-project", + location="global", + engine_id="test-engine", + search_query="Google Cloud revenue", + data_store_ids=["chunked-ds", "unchunked-ds"], + ) + ) + + self.assertEqual(response, mock_response) + mock_client.search.assert_called_once() + call_args = mock_client.search.call_args[1]["request"] + + self.assertIn("engines/test-engine", call_args.serving_config) + self.assertEqual(call_args.query, "Google Cloud revenue") + extractive_spec = call_args.content_search_spec.extractive_content_spec + self.assertEqual(extractive_spec.max_extractive_segment_count, 3) + self.assertTrue(extractive_spec.return_extractive_segment_score) + self.assertEqual(len(call_args.data_store_specs), 2) + + @mock.patch("google.cloud.discoveryengine_v1.UserEventServiceClient") + def test_write_user_event_with_attribution_with_data_store( + self, mock_client_cls + ): + mock_client = mock.MagicMock() + mock_client_cls.return_value = mock_client + + mock_doc_info = discoveryengine.DocumentInfo( + name="projects/test-project/locations/global/collections/default_collection/dataStores/target-ds/branches/0/documents/doc_456" + ) + mock_response = discoveryengine.UserEvent( + event_type="view-item", + user_pseudo_id="visitor_999", + engine="projects/test-project/locations/global/collections/default_collection/engines/test-engine", + data_store="projects/test-project/locations/global/collections/default_collection/dataStores/target-ds", + attribution_token="search_token_abc", + documents=[mock_doc_info], + ) + mock_client.write_user_event.return_value = mock_response + + response = write_user_event_sample.write_user_event_with_attribution( + project_id="test-project", + location="global", + engine_id="test-engine", + user_pseudo_id="visitor_999", + attribution_token="search_token_abc", + document_id="doc_456", + data_store_id="target-ds", + ) + + self.assertEqual(response, mock_response) + mock_client.write_user_event.assert_called_once() + req = mock_client.write_user_event.call_args[1]["request"] + + self.assertEqual(req.parent, "projects/test-project/locations/global") + self.assertEqual(req.user_event.event_type, "view-item") + self.assertEqual(req.user_event.attribution_token, "search_token_abc") + self.assertEqual(req.user_event.user_pseudo_id, "visitor_999") + self.assertIn("documents/doc_456", req.user_event.documents[0].name) + self.assertIn("engines/test-engine", req.user_event.engine) + self.assertIn("dataStores/target-ds", req.user_event.data_store) + + @mock.patch("google.cloud.discoveryengine_v1.UserEventServiceClient") + def test_write_user_event_with_attribution_no_data_store( + self, mock_client_cls + ): + mock_client = mock.MagicMock() + mock_client_cls.return_value = mock_client + + mock_doc_info = discoveryengine.DocumentInfo(id="doc_789") + mock_response = discoveryengine.UserEvent( + event_type="view-item", + user_pseudo_id="visitor_111", + engine="projects/test-project/locations/global/collections/default_collection/engines/test-engine", + attribution_token="search_token_xyz", + documents=[mock_doc_info], + ) + mock_client.write_user_event.return_value = mock_response + + response = write_user_event_sample.write_user_event_with_attribution( + project_id="test-project", + location="global", + engine_id="test-engine", + user_pseudo_id="visitor_111", + attribution_token="search_token_xyz", + document_id="doc_789", + ) + + self.assertEqual(response, mock_response) + mock_client.write_user_event.assert_called_once() + req = mock_client.write_user_event.call_args[1]["request"] + + self.assertEqual(req.parent, "projects/test-project/locations/global") + self.assertEqual(req.user_event.documents[0].id, "doc_789") + self.assertFalse(req.user_event.data_store) + + @mock.patch( + "google.cloud.discoveryengine_v1.ConversationalSearchServiceClient" + ) + def test_converse_conversation_new_session(self, mock_client_cls): + mock_client = mock.MagicMock() + mock_client_cls.return_value = mock_client + + mock_created_conv = discoveryengine.Conversation( + name="projects/test-project/locations/global/collections/default_collection/dataStores/test-ds/conversations/conv-123" + ) + mock_client.create_conversation.return_value = mock_created_conv + + mock_summary = discoveryengine.SearchResponse.Summary( + summary_text="Vertex AI Search provides enterprise generative search." + ) + mock_reply = discoveryengine.Reply(summary=mock_summary) + mock_search_result = discoveryengine.SearchResponse.SearchResult( + id="doc-1", + document=discoveryengine.Document( + id="doc-1", + struct_data={ + "title": "VAIS Docs", + "link": "https://cloud.google.com", + }, + ), + ) + mock_response = discoveryengine.ConverseConversationResponse( + reply=mock_reply, + conversation=mock_created_conv, + search_results=[mock_search_result], + ) + mock_client.converse_conversation.return_value = mock_response + + response = converse_conversation_sample.multi_turn_conversational_search( + project_id="test-project", + location="global", + data_store_id="test-ds", + query_text="What is Vertex AI Search?", + ) + + self.assertEqual(response, mock_response) + mock_client.create_conversation.assert_called_once() + mock_client.converse_conversation.assert_called_once() + + @mock.patch( + "google.cloud.discoveryengine_v1.ConversationalSearchServiceClient" + ) + def test_converse_conversation_existing_session(self, mock_client_cls): + mock_client = mock.MagicMock() + mock_client_cls.return_value = mock_client + + mock_summary = discoveryengine.SearchResponse.Summary( + summary_text="It also supports grounded citation parsing." + ) + mock_reply = discoveryengine.Reply(summary=mock_summary) + mock_response = discoveryengine.ConverseConversationResponse( + reply=mock_reply, + search_results=[], + ) + mock_client.converse_conversation.return_value = mock_response + + response = converse_conversation_sample.multi_turn_conversational_search( + project_id="test-project", + location="global", + data_store_id="test-ds", + query_text="Tell me about citations.", + conversation_id="conv-existing-999", + ) + + self.assertEqual(response, mock_response) + mock_client.create_conversation.assert_not_called() + mock_client.converse_conversation.assert_called_once() + req = mock_client.converse_conversation.call_args[1]["request"] + self.assertIn("conv-existing-999", req.name) + + @mock.patch("google.cloud.discoveryengine_v1.DocumentServiceClient") + def test_create_structured_document_with_metadata(self, mock_client_cls): + mock_client = mock.MagicMock() + mock_client_cls.return_value = mock_client + + mock_doc = discoveryengine.Document( + id="doc-storage-01", + name="projects/test-project/locations/global/collections/default_collection/dataStores/test-ds/branches/0/documents/doc-storage-01", + struct_data={ + "title": "Cloud Storage Architecture", + "url": "https://example.com/docs/gcs_storage_guide_v2", + "category": "Storage", + "rating": 4.9, + "tags": ["gcs", "cloud", "infra"], + }, + ) + mock_client.create_document.return_value = mock_doc + + response = ( + create_document_metadata_sample.create_structured_document_with_metadata( + project_id="test-project", + location="global", + data_store_id="test-ds", + document_id="doc-storage-01", + title="Cloud Storage Architecture", + uri="https://example.com/docs/gcs_storage_guide_v2", + category="Storage", + rating=4.9, + tags=["gcs", "cloud", "infra"], + content_text=( + "Overview of Google Cloud Storage buckets and security." + ), + ) + ) + + self.assertEqual(response, mock_doc) + mock_client.create_document.assert_called_once() + req = mock_client.create_document.call_args[1]["request"] + self.assertIn("branches/0", req.parent) + self.assertEqual(req.document_id, "doc-storage-01") + self.assertEqual( + req.document.struct_data.get("url"), + "https://example.com/docs/gcs_storage_guide_v2", + ) + self.assertEqual(req.document.content.mime_type, "text/plain") + + @mock.patch("google.cloud.discoveryengine_v1.DocumentServiceClient") + def test_poll_long_running_operation_robust_success(self, mock_client_cls): + mock_client = mock.MagicMock() + mock_client_cls.return_value = mock_client + + mock_op_running = operations_pb2.Operation( + name="projects/p/locations/l/operations/op-123", + done=False, + ) + mock_metadata = any_pb2.Any( + type_url="type.googleapis.com/google.cloud.discoveryengine.v1.ImportDocumentsMetadata" + ) + mock_op_done = operations_pb2.Operation( + name="projects/p/locations/l/operations/op-123", + done=True, + metadata=mock_metadata, + ) + mock_client.transport.operations_client.get_operation.side_effect = [ + mock_op_running, + mock_op_done, + ] + + with mock.patch("time.sleep") as mock_sleep: + op = poll_lro_robust_sample.poll_long_running_operation_robust( + operation_name="projects/p/locations/l/operations/op-123", + initial_delay_seconds=1.0, + max_delay_seconds=5.0, + timeout_seconds=30.0, + ) + self.assertTrue(op.done) + mock_sleep.assert_called_once() + + @mock.patch("google.cloud.discoveryengine_v1.DocumentServiceClient") + def test_poll_long_running_operation_robust_error(self, mock_client_cls): + mock_client = mock.MagicMock() + mock_client_cls.return_value = mock_client + + mock_op_failed = operations_pb2.Operation( + name="projects/p/locations/l/operations/op-err", + done=True, + error=status_pb2.Status(code=3, message="Invalid GCS bucket uri"), + ) + mock_client.transport.operations_client.get_operation.return_value = ( + mock_op_failed + ) + + with self.assertRaises(RuntimeError) as ctx: + poll_lro_robust_sample.poll_long_running_operation_robust( + operation_name="projects/p/locations/l/operations/op-err", + timeout_seconds=30.0, + ) + self.assertIn("Invalid GCS bucket uri", str(ctx.exception)) + + @mock.patch("google.cloud.discoveryengine_v1.DocumentServiceClient") + def test_poll_long_running_operation_robust_timeout(self, mock_client_cls): + mock_client = mock.MagicMock() + mock_client_cls.return_value = mock_client + + mock_op_running = operations_pb2.Operation( + name="projects/p/locations/l/operations/op-timeout", + done=False, + ) + mock_client.transport.operations_client.get_operation.return_value = ( + mock_op_running + ) + + with mock.patch("time.time") as mock_time, mock.patch("time.sleep"): + mock_time.side_effect = [0.0, 0.0, 10.0, 50.0] + with self.assertRaises(TimeoutError): + poll_lro_robust_sample.poll_long_running_operation_robust( + operation_name="projects/p/locations/l/operations/op-timeout", + timeout_seconds=20.0, + ) + + def test_create_data_store_with_gemini_parser(self): + de = enable_gemini_layout_parser_sample.discoveryengine + with mock.patch.object(de, "DataStoreServiceClient") as mock_client_cls: + mock_client = mock.MagicMock() + mock_client_cls.return_value = mock_client + + mock_ds = de.DataStore( + name="projects/test-project/locations/global/collections/default_collection/dataStores/gemini-ds", + display_name="Gemini Parsed DS", + ) + mock_operation = mock.MagicMock() + mock_operation.operation.name = ( + "projects/test-project/locations/global/operations/op-ds-create" + ) + mock_operation.result.return_value = mock_ds + mock_client.create_data_store.return_value = mock_operation + + response = enable_gemini_layout_parser_sample.create_data_store_with_gemini_parser( + project_id="test-project", + location="global", + data_store_id="gemini-ds", + display_name="Gemini Parsed DS", + enable_table_annotation=True, + enable_image_annotation=True, + ) + + self.assertEqual(response, mock_ds) + mock_client.create_data_store.assert_called_once() + req = mock_client.create_data_store.call_args[1]["request"] + self.assertEqual(req.data_store_id, "gemini-ds") + self.assertEqual(req.data_store.display_name, "Gemini Parsed DS") + + layout_cfg = ( + req.data_store.document_processing_config.default_parsing_config.layout_parsing_config + ) + self.assertTrue(layout_cfg.enable_llm_layout_parsing) + self.assertTrue(layout_cfg.enable_table_annotation) + self.assertTrue(layout_cfg.enable_image_annotation) + + chunking_cfg = ( + req.data_store.document_processing_config.chunking_config.layout_based_chunking_config + ) + self.assertEqual(chunking_cfg.chunk_size, 500) + self.assertTrue(chunking_cfg.include_ancestor_headings) + + def test_update_data_store_gemini_parser(self): + de = update_data_store_gemini_parser_sample.discoveryengine + with mock.patch.object(de, "DataStoreServiceClient") as mock_client_cls: + mock_client = mock.MagicMock() + mock_client_cls.return_value = mock_client + + mock_ds = de.DataStore( + name="projects/test-project/locations/global/collections/default_collection/dataStores/gemini-ds", + display_name="Gemini Parsed DS", + ) + mock_client.update_data_store.return_value = mock_ds + + response = ( + update_data_store_gemini_parser_sample.update_data_store_gemini_parser( + project_id="test-project", + location="global", + data_store_id="gemini-ds", + enable_table_annotation=True, + enable_image_annotation=True, + ) + ) + + self.assertEqual(response, mock_ds) + mock_client.update_data_store.assert_called_once() + req = mock_client.update_data_store.call_args[1]["request"] + self.assertIn("dataStores/gemini-ds", req.data_store.name) + self.assertEqual( + list(req.update_mask.paths), ["document_processing_config"] + ) + layout_cfg = ( + req.data_store.document_processing_config.default_parsing_config.layout_parsing_config + ) + self.assertTrue(layout_cfg.enable_llm_layout_parsing) + self.assertTrue(layout_cfg.enable_table_annotation) + self.assertTrue(layout_cfg.enable_image_annotation) + + def test_import_custom_chunk_documents(self): + de = import_custom_chunks_sample.discoveryengine + with mock.patch.object(de, "DocumentServiceClient") as mock_client_cls: + mock_client = mock.MagicMock() + mock_client_cls.return_value = mock_client + + mock_op = mock.MagicMock() + mock_op.operation.name = "projects/test-project/locations/global/collections/default_collection/dataStores/test-ds/branches/0/operations/op-byoc-123" + mock_client.import_documents.return_value = mock_op + + response = import_custom_chunks_sample.import_custom_chunk_documents( + project_id="test-project", + location="global", + data_store_id="test-ds", + gcs_uri="gs://my-bucket/custom_chunks/*.json", + ) + + self.assertEqual(response, mock_op) + mock_client.import_documents.assert_called_once() + req = mock_client.import_documents.call_args[1]["request"] + self.assertIn("dataStores/test-ds/branches/0", req.parent) + self.assertEqual( + req.gcs_source.input_uris, ["gs://my-bucket/custom_chunks/*.json"] + ) + self.assertEqual(req.gcs_source.data_schema, "custom") + self.assertEqual( + req.reconciliation_mode, + de.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL, + ) + + def test_search_chunks_with_metadata(self): + de = search_chunks_sample.discoveryengine + with mock.patch.object(de, "SearchServiceClient") as mock_client_cls: + mock_client = mock.MagicMock() + mock_client_cls.return_value = mock_client + + doc_metadata = de.Chunk.DocumentMetadata( + title="Discovery Engine Deep Dive", + uri="https://cloud.google.com/generative-ai-app-builder/docs/parse-chunk-documents", + struct_data={"category": "Enterprise Search", "rating": 4.9}, + ) + + prev_chunk = de.Chunk(id="c1", content="Preceding context on indexing.") + next_chunk = de.Chunk(id="c3", content="Succeeding context on ranking.") + + chunk_meta = de.Chunk.ChunkMetadata( + previous_chunks=[prev_chunk], + next_chunks=[next_chunk], + ) + + mock_chunk = de.Chunk( + id="c2", + name="projects/test-p/locations/global/collections/default_collection/dataStores/chunk-ds/branches/0/documents/doc1/chunks/c2", + content=( + "Main relevant chunk discussing layout-aware document chunking." + ), + relevance_score=0.965, + document_metadata=doc_metadata, + page_span=de.Chunk.PageSpan(page_start=5, page_end=6), + chunk_metadata=chunk_meta, + ) + + mock_search_result = de.SearchResponse.SearchResult( + id="doc1", + chunk=mock_chunk, + ) + + mock_response = de.SearchResponse( + results=[mock_search_result], + attribution_token="test_attribution_token_chunk_999", + total_size=1, + ) + mock_client.search.return_value = mock_response + + response = search_chunks_sample.search_chunks_with_metadata( + project_id="test-p", + location="global", + data_store_id="chunk-ds", + search_query="how to use chunking", + num_previous_chunks=1, + num_next_chunks=1, + filter_expr='category: ANY("Enterprise Search")', + ) + + self.assertEqual(response, mock_response) + mock_client.search.assert_called_once() + req = mock_client.search.call_args[1]["request"] + self.assertEqual(req.query, "how to use chunking") + self.assertEqual(req.filter, 'category: ANY("Enterprise Search")') + self.assertEqual( + req.content_search_spec.search_result_mode, + de.SearchRequest.ContentSearchSpec.SearchResultMode.CHUNKS, + ) + self.assertEqual( + req.content_search_spec.chunk_spec.num_previous_chunks, 1 + ) + self.assertEqual(req.content_search_spec.chunk_spec.num_next_chunks, 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/discoveryengine/update_data_store_gemini_parser_sample.py b/discoveryengine/update_data_store_gemini_parser_sample.py new file mode 100644 index 00000000000..b91278a69aa --- /dev/null +++ b/discoveryengine/update_data_store_gemini_parser_sample.py @@ -0,0 +1,97 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Gemini Advanced Layout Parser update sample for Agent Search.""" + +# [START genappbuilder_update_data_store_gemini_parser] +from google.api_core.client_options import ClientOptions +from google.cloud import discoveryengine_v1beta as discoveryengine +from google.protobuf import field_mask_pb2 + +FieldMask = field_mask_pb2.FieldMask + + +def update_data_store_gemini_parser( + project_id: str, + location: str, + data_store_id: str, + enable_table_annotation: bool = True, + enable_image_annotation: bool = True, +) -> discoveryengine.DataStore: + """Updates an existing DataStore to enable Gemini Advanced Layout Parsing. + + Updating DocumentProcessingConfig modifies how subsequent document imports + are parsed and chunked. Previously ingested documents are not automatically + re-parsed; to apply Gemini layout parsing to existing documents, re-import + them after updating the config. + + Args: + project_id: Google Cloud project ID. + location: DataStore location (e.g., 'global', 'us', 'eu'). + data_store_id: Unique identifier of the existing DataStore. + enable_table_annotation: Whether to generate LLM descriptions for tables. + enable_image_annotation: Whether to generate LLM descriptions for images. + + Returns: + The updated DataStore resource. + """ + client_options = ( + ClientOptions(api_endpoint=f"{location}-discoveryengine.googleapis.com") + if location != "global" + else None + ) + client = discoveryengine.DataStoreServiceClient(client_options=client_options) + + data_store_name = ( + f"projects/{project_id}/locations/{location}/collections/default_collection/" + f"dataStores/{data_store_id}" + ) + + # 1. Build LayoutParsingConfig with Gemini enhancement + layout_parsing_config = discoveryengine.DocumentProcessingConfig.ParsingConfig.LayoutParsingConfig( + enable_llm_layout_parsing=True, + enable_table_annotation=enable_table_annotation, + enable_image_annotation=enable_image_annotation, + ) + + parsing_config = discoveryengine.DocumentProcessingConfig.ParsingConfig( + layout_parsing_config=layout_parsing_config + ) + + # 2. Build updated DocumentProcessingConfig + doc_processing_config = discoveryengine.DocumentProcessingConfig( + default_parsing_config=parsing_config, + ) + + # 3. Construct DataStore with field mask for document_processing_config + data_store = discoveryengine.DataStore( + name=data_store_name, + document_processing_config=doc_processing_config, + ) + field_mask = FieldMask(paths=["document_processing_config"]) + + request = discoveryengine.UpdateDataStoreRequest( + data_store=data_store, + update_mask=field_mask, + ) + + updated_data_store = client.update_data_store(request=request) + print("Successfully updated DataStore to Gemini Parser:") + print(f" Name: {updated_data_store.name}") + print(" Default Parser: LayoutParser (Gemini LLM-enhanced)") + + return updated_data_store + + +# [END genappbuilder_update_data_store_gemini_parser] diff --git a/discoveryengine/write_user_event_sample.py b/discoveryengine/write_user_event_sample.py new file mode 100644 index 00000000000..53c2800da5e --- /dev/null +++ b/discoveryengine/write_user_event_sample.py @@ -0,0 +1,129 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""User event recording with search attribution sample for Agent Search.""" + +# [START genappbuilder_write_user_event] +import time +from typing import Optional + +from google.cloud import discoveryengine_v1 as discoveryengine +from google.protobuf import timestamp_pb2 + + +def write_user_event_with_attribution( + project_id: str, + location: str, + engine_id: str, + user_pseudo_id: str, + attribution_token: str, + document_id: str, + event_type: str = "view-item", + data_store_id: Optional[str] = None, +) -> discoveryengine.UserEvent: + """Records a real-time user event scoped to an engine with search attribution. + + To ensure analytics pipelines correctly attribute click events with preceding + search queries and maintain accurate Click-Through Rate (CTR) analytics: + 1. Ingest events at the Location level: + `projects/{project}/locations/{location}` + 2. Populate the `engine` resource field on the `UserEvent` proto. + 3. Pass the exact `attribution_token` from the corresponding `SearchResponse`. + 4. Populate the `documents` list with the clicked document ID. + + Args: + project_id: Google Cloud project ID or project number. + location: Engine location (e.g., 'global', 'us', 'eu'). + engine_id: Search engine (app) ID. + user_pseudo_id: Uniquely pseudonymized visitor/session identifier. + attribution_token: Token received from the preceding search response. + document_id: ID of the document viewed or clicked. + event_type: Event type (e.g., 'view-item', 'search', 'view-category'). + data_store_id: Optional specific data store ID associated with the doc. + + Returns: + The recorded UserEvent. + """ + client = discoveryengine.UserEventServiceClient() + + # CRITICAL: Always use the Location-level parent for Engine-scoped events + parent = f"projects/{project_id}/locations/{location}" + + engine_path = ( + f"projects/{project_id}/locations/{location}/collections/default_collection" + f"/engines/{engine_id}" + ) + + # Set accurate UTC event timestamp + current_time = time.time() + event_time = timestamp_pb2.Timestamp( + seconds=int(current_time), + nanos=int((current_time - int(current_time)) * 1e9), + ) + + # Build DocumentInfo reference for the clicked/viewed item + if document_id.startswith("projects/"): + document_info = discoveryengine.DocumentInfo(name=document_id) + elif data_store_id: + doc_name = ( + f"projects/{project_id}/locations/{location}/collections/default_collection" + f"/dataStores/{data_store_id}/branches/0/documents/{document_id}" + ) + document_info = discoveryengine.DocumentInfo(name=doc_name) + else: + document_info = discoveryengine.DocumentInfo(id=document_id) + + user_event = discoveryengine.UserEvent( + event_type=event_type, + user_pseudo_id=user_pseudo_id, + engine=engine_path, + attribution_token=attribution_token, + documents=[document_info], + event_time=event_time, + user_info=discoveryengine.UserInfo( + user_id=f"user_{user_pseudo_id}", + user_agent=( + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36" + ), + ), + page_info=discoveryengine.PageInfo( + uri="https://example.com/search/results", + pageview_id="pageview_12345", + ), + ) + + if data_store_id: + user_event.data_store = ( + f"projects/{project_id}/locations/{location}/collections/default_collection" + f"/dataStores/{data_store_id}" + ) + + request = discoveryengine.WriteUserEventRequest( + parent=parent, + user_event=user_event, + ) + + response = client.write_user_event(request=request) + + print(f"Recorded User Event '{response.event_type}':") + print(f" User Pseudo ID: {response.user_pseudo_id}") + print(f" Engine: {response.engine}") + print(f" Attribution Token: {response.attribution_token}") + print(f" Documents: {[doc.id or doc.name for doc in response.documents]}") + print(f" Event Time: {response.event_time}") + + return response + + +# [END genappbuilder_write_user_event] From d776f6587b40319616b2d4fc2a0d82ad55688e91 Mon Sep 17 00:00:00 2001 From: sailobo Date: Tue, 8 Sep 2026 20:32:55 -0600 Subject: [PATCH 2/4] Update write_user_event_sample.py --- discoveryengine/write_user_event_sample.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/discoveryengine/write_user_event_sample.py b/discoveryengine/write_user_event_sample.py index 53c2800da5e..81961324c2f 100644 --- a/discoveryengine/write_user_event_sample.py +++ b/discoveryengine/write_user_event_sample.py @@ -66,12 +66,8 @@ def write_user_event_with_attribution( ) # Set accurate UTC event timestamp - current_time = time.time() - event_time = timestamp_pb2.Timestamp( - seconds=int(current_time), - nanos=int((current_time - int(current_time)) * 1e9), - ) - + event_time = timestamp_pb2.Timestamp() + event_time.GetCurrentTime() # Build DocumentInfo reference for the clicked/viewed item if document_id.startswith("projects/"): document_info = discoveryengine.DocumentInfo(name=document_id) From 9636fa0c854278f80a2131766f43148a8a7f99ff Mon Sep 17 00:00:00 2001 From: sailobo Date: Tue, 8 Sep 2026 20:37:04 -0600 Subject: [PATCH 3/4] Update converse_conversation_sample.py --- discoveryengine/converse_conversation_sample.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/discoveryengine/converse_conversation_sample.py b/discoveryengine/converse_conversation_sample.py index 4bdffe01258..60c3faa1add 100644 --- a/discoveryengine/converse_conversation_sample.py +++ b/discoveryengine/converse_conversation_sample.py @@ -95,15 +95,12 @@ def multi_turn_conversational_search( response = client.converse_conversation(request=request) reply_text = "" + if response.reply: - if ( - hasattr(response.reply, "summary") - and response.reply.summary - and response.reply.summary.summary_text - ): - reply_text = response.reply.summary.summary_text - elif hasattr(response.reply, "reply") and response.reply.reply: - reply_text = response.reply.reply + if response.reply.summary and response.reply.summary.summary_text: + reply_text = response.reply.summary.summary_text + elif response.reply.reply: + reply_text = response.reply.reply print(f"\nUser Query: {query_text}") print(f"AI Generated Reply: {reply_text}") From c72fc88e608ee543beb98656bce7e9738bb5576c Mon Sep 17 00:00:00 2001 From: sailobo Date: Tue, 8 Sep 2026 20:48:49 -0600 Subject: [PATCH 4/4] Apply batched suggestions from code review Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- discoveryengine/create_document_metadata_sample.py | 10 ++++++---- discoveryengine/search_extractive_segments_sample.py | 3 +++ 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/discoveryengine/create_document_metadata_sample.py b/discoveryengine/create_document_metadata_sample.py index ec2fb1cb617..866bca59399 100644 --- a/discoveryengine/create_document_metadata_sample.py +++ b/discoveryengine/create_document_metadata_sample.py @@ -100,10 +100,12 @@ def create_structured_document_with_metadata( print(f"Created Document ID: {response.id}") print(f" Name: {response.name}") - print(f" Metadata URL (exact): {response.struct_data.get('url')}") - print(f" Category: {response.struct_data.get('category')}") - print(f" Rating: {response.struct_data.get('rating')}") - print(f" Tags: {response.struct_data.get('tags')}") + struct_data = response.struct_data + if struct_data: + print(f" Metadata URL (exact): {struct_data.get('url')}") + print(f" Category: {struct_data.get('category')}") + print(f" Rating: {struct_data.get('rating')}") + print(f" Tags: {struct_data.get('tags')}") return response diff --git a/discoveryengine/search_extractive_segments_sample.py b/discoveryengine/search_extractive_segments_sample.py index 2e76d908d7e..821e8fd54b4 100644 --- a/discoveryengine/search_extractive_segments_sample.py +++ b/discoveryengine/search_extractive_segments_sample.py @@ -94,6 +94,9 @@ def search_with_extractive_segments( for result in response.results: doc = result.document + if not doc: + print("\n- (No document returned)") + continue print(f"\n- Document ID: {doc.id}") print(f" Name: {doc.name}")