Add code samples for Agent Search

discoveryengine/converse_conversation_sample.py

@@ -0,0 +1,123 @@
+# 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 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}")
+
+  # 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]

discoveryengine/create_document_metadata_sample.py

@@ -0,0 +1,113 @@
+# 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}")
+  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
+
+
+# [END genappbuilder_create_document_metadata]

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]