wehub/googlecloudplatform--generative-ai
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:30:30 +08:00
commit 914fea506e
2793 changed files with 802106 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
search/cloud-function/README.md
+128
View File
@@ -0,0 +1,128 @@
# Vertex AI Search accessed via Google Cloud Functions
This directory contains several versions of approximately the same
implementation.
The functions can be deployed to
[Cloud functions](https://cloud.google.com/functions/) and can be modified to
supports many different triggers and use cases. Each can also
[be deployed locally](https://cloud.google.com/functions/docs/running/overview)
which allows easy experimentation and iteration.
This example is powered by
[Vertex AI Search](https://cloud.google.com/generative-ai-app-builder/docs/enterprise-search-introduction)
which does many different things, including **Document & Intranet Search**,
**Recommendations** and **Grounding and RAG** out-of-the-box (For more
information, see the blog post
[Your RAG powered by Google Search](https://cloud.google.com/blog/products/ai-machine-learning/rags-powered-by-google-search-technology-part-1)).
If you want even more control see
[Vertex AI Search Component APIs](https://cloud.google.com/generative-ai-app-builder/docs/builder-apis),
but first explore the out-of-the-box offering because it's easy to setup.
This example is for the out-of-the-box Vertex AI Search supporting many
configurations and data types.
## Pre-requisites
Before you can use these functions to query Vertex AI Search, you need to create
and populate a search "data store"; read through instructions in
[get started with generic search](https://cloud.google.com/generative-ai-app-builder/docs/try-enterprise-search).
These functions could easily be adapted to other types of Vertex AI Search like
[generic recommendations](https://cloud.google.com/generative-ai-app-builder/docs/try-generic-recommendations),
[media search](https://cloud.google.com/generative-ai-app-builder/docs/try-media-search),
[media recommendations](https://cloud.google.com/generative-ai-app-builder/docs/try-media-recommendations),
[healthcare search](https://cloud.google.com/generative-ai-app-builder/docs/create-data-store-hc),
or even
[retail product discovery](https://cloud.google.com/solutions/retail-product-discovery#documentation).
You'll need to collect the following details from your search app data store:
```python
PROJECT_ID = "YOUR_PROJECT_ID" # alphanumeric
LOCATION = "global" # or an alternate location
DATA_STORE_ID = "YOUR_DATA_STORE_ID" # not the app id, alphanumeric
```
Additionally you'll need to keep track of some of the choices you make when you
configure Vertex AI Search.
### Type of data source
<!-- textlint-disable -->
- UNSTRUCTURED
- STRUCTURED
- WEBSITE
- BLENDED
<!-- textlint-enable -->
```python
ENGINE_DATA_TYPE = UNSTRUCTURED
```
### Type of chunks to return
- DOCUMENT_WITH_SNIPPETS
- DOCUMENT_WITH_EXTRACTIVE_SEGMENTS
- CHUNK
- NONE
```python
ENGINE_CHUNK_TYPE = DOCUMENT_WITH_EXTRACTIVE_SEGMENTS
```
### Type of summarization
- NONE results only
- VERTEX_AI_SEARCH LLM add on provided by
[Vertex AI Search](https://cloud.google.com/generative-ai-app-builder/docs/enterprise-search-introduction)
<!-- NOT ready yet
- GENERATE_GROUNDED_ANSWERS use the
[Generate grounded answers with RAG](https://cloud.google.com/generative-ai-app-builder/docs/grounded-gen)
provided by
[Vertex AI Search Builder APIs](https://cloud.google.com/generative-ai-app-builder/docs/builder-apis)
- GEMINI use one of the Gemini models to generate an answer from the results -->
```python
SUMMARY_TYPE = VERTEX_AI_SEARCH
```
## Architecture
1. Vertex AI Search is an API hosted on Google Cloud
2. You will call that API via a Google Cloud Function, which exposes its own API
3. Your users will the Google Cloud Function API, via your custom app or UI
```mermaid
flowchart LR
A[fa:fa-search Vertex AI Search] --> B(Google Cloud Function)
B --> C[My App Server]
C -->|One| D[fa:fa-laptop web]
C -->|Two| E[fa:fa-mobile mobile]
```
## Use case: RAG / Grounding
Any time you have more source data than can fit into a LLM context window, you
could benefit from RAG (Retrieval Augmented Generation). The more data you have,
the more important search is - to get the relevant chunks into the prompt of the
LLM.
- **Retrieve** relevant search results, with text chunks (snippets or segments)
- **Augmented Generation** uses Gemini to generate an answer or summary grounded
on the relevant search results
## Use case: Agent Tool (Knowledge Base)
A natural extension of RAG / Grounding is agentic behavior.
Whether creating a basic chatbot or a sophisticated tool using multi-agent
system, you're always going to need search based RAG. The better the search
quality the better the agent response based on your source data.
For more on agents, check out
[Vertex AI Search Use Cases](https://cloud.google.com/products/agent-builder?hl=en#common-uses)
and
[https://github.com/GoogleCloudPlatform/generative-ai](https://github.com/GoogleCloudPlatform/generative-ai).
search/cloud-function/python/.env.example
+6
View File
@@ -0,0 +1,6 @@
DATA_STORE_ID=your-data-store-id-here
ENGINE_CHUNK_TYPE=DOCUMENT_WITH_EXTRACTIVE_SEGMENTS
ENGINE_DATA_TYPE=UNSTRUCTURED
LOCATION=global
PROJECT_ID=your-project-id-here
SUMMARY_TYPE=VERTEX_AI_SEARCH
search/cloud-function/python/README.md
+127
View File
@@ -0,0 +1,127 @@
# Vertex AI Search accessed via Google Cloud Functions
This example is based on the
[Python client for the Vertex AI Search API](https://cloud.google.com/generative-ai-app-builder/docs/libraries#client-libraries-usage-python),
which will get search results, snippets, metadata, and the LLM summary grounded
on search results. This is implemented in the `vertex_ai_search_client.py` file.
That functionality is exposed on a REST API which is implemented in `main.py`
intended to be deployed to a Google Cloud Function using an HTTPS trigger on a
Python 3 runtime;
[read more here](https://cloud.google.com/functions/docs/samples/functions-http-content#functions_http_content-python).
**[Read more about Vertex AI Search accessed via Google Cloud Functions](../)**
## Environment Variables
The following environment variables are required for both local development and
deployment:
- `PROJECT_ID`: Your Google Cloud project ID
- `LOCATION`: The location of your Vertex AI Search data store
- `DATA_STORE_ID`: The ID of your Vertex AI Search data store
- `ENGINE_DATA_TYPE`: Type of data in the engine (0-3)
- `ENGINE_CHUNK_TYPE`: Type of chunking used (0-3)
- `SUMMARY_TYPE`: Type of summary used (0-3)
## Local Development
### Setup
1. Ensure you have the Google Cloud SDK installed and configured.
2. Clone this repository and navigate to the project directory.
3. Set up your environment variables:
```bash
gcloud auth login
bash setup_env.sh
```
Alternatively, you can manually create and edit a `.env` file with the required
variables.
### Run locally
Run this code locally via **Functions Framework** or **Functions Emulator**;
[read more about running cloud functions locally](https://cloud.google.com/functions/docs/running/overview).
```bash
pip install -r requirements.txt
pip install functions-framework
functions-framework --target=vertex_ai_search
```
In a different terminal, execute a `POST` search query based on your data:
```bash
export SEARCH_TERM="What is the ... for ...?"
curl -m 310 -X POST localhost:8080 \
-H "Content-Type: application/json" \
-d "{\"search_term\": \"${SEARCH_TERM}\"}"
```
### Run tests
#### Unit tests
These tests mock the API interactions and should run quickly:
```bash
pip install pytest
pytest test_vertex_ai_search_client.py
```
#### Integration tests
These tests actually call the Vertex AI Search API and depend on your data
stores being configured in Vertex AI Search:
```bash
pip install pytest
pytest test_integration_vertex_ai_search_client.py
```
## Deployment
To deploy this function to Google Cloud:
1. Ensure you have set up the required environment variables (see Environment
Variables section).
2. Run the following command:
```bash
gcloud functions deploy vertex_ai_search --runtime python39 --trigger-http --allow-unauthenticated
```
You will get back a URL for triggering the function.
## Usage
After deployment, you can use the function as follows:
```bash
curl -X POST https://YOUR_FUNCTION_URL \
-H "Content-Type: application/json" \
-d '{"search_term": "your search query"}'
```
Replace `YOUR_FUNCTION_URL` with the URL of your deployed function, and fill in
the search query.
If you run into problems, go to
[Google Cloud Functions](https://console.cloud.google.com/functions), find the
function you just deployed, and review the logs for informative errors. Perhaps
you need to setup
[Google Cloud IAM](https://cloud.google.com/functions/docs/reference/iam) roles
or permissions.
## Customization
This implementation provides a basic way to access and control your queries to
the Vertex AI Search API. It simplifies CORS and bearer token authentication,
and allows for some minor customization of inputs and outputs.
If you require more extensive customization, consider using an orchestration
framework like [LangChain](https://www.langchain.com/) or
[LlamaIndex](https://www.llamaindex.ai/) which have Vertex AI Search
integrations.
search/cloud-function/python/main.py
+135
View File
@@ -0,0 +1,135 @@
# Copyright 2024 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
#
# http://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.
"""
Google Cloud Function for Vertex AI Search
This module provides an HTTP endpoint for performing searches using
the Vertex AI Search API. It uses the VertexAISearchClient to handle
the core search functionality.
For deployment instructions, environment variable setup, and usage examples,
please refer to the README.md file.
"""
import os
from typing import Any
from flask import Flask, Request, jsonify, request
import functions_framework
from google.api_core.exceptions import GoogleAPICallError
from vertex_ai_search_client import VertexAISearchClient, VertexAISearchConfig
# Load environment variables
project_id = os.getenv("PROJECT_ID", "your-project")
location = os.getenv("LOCATION", "global")
data_store_id = os.getenv("DATA_STORE_ID", "your-data-store")
engine_data_type = os.getenv("ENGINE_DATA_TYPE", "UNSTRUCTURED")
engine_chunk_type = os.getenv("ENGINE_CHUNK_TYPE", "CHUNK")
summary_type = os.getenv("SUMMARY_TYPE", "VERTEX_AI_SEARCH")
# Create VertexAISearchConfig
config = VertexAISearchConfig(
project_id=project_id,
location=location,
data_store_id=data_store_id,
engine_data_type=engine_data_type,
engine_chunk_type=engine_chunk_type,
summary_type=summary_type,
)
# Initialize VertexAISearchClient
vertex_ai_search_client = VertexAISearchClient(config)
@functions_framework.http
def vertex_ai_search(http_request: Request) -> tuple[Any, int, dict[str, str]]:
"""
Handle HTTP requests for Vertex AI Search.
This function processes incoming HTTP requests, performs the search using
the VertexAISearchClient, and returns the results. It handles CORS, validates
the request, and manages potential errors.
Args:
http_request (flask.Request): The request object.
<https://flask.palletsprojects.com/en/1.1.x/api/#incoming-request-data>
Returns:
Tuple[Any, int, Dict[str, str]]: A tuple containing the response body,
status code, and headers. This output will be turned into a Response
object using `make_response`
<https://flask.palletsprojects.com/en/1.1.x/api/#flask.make_response>.
"""
# Set CORS headers for the preflight request
if http_request.method == "OPTIONS":
# Allows GET requests from any origin with the Content-Type
# header and caches preflight response for an 3600s
headers = {
"Access-Control-Allow-Origin": "*",
"Access-Control-Allow-Methods": "GET",
"Access-Control-Allow-Headers": "Content-Type",
"Access-Control-Max-Age": "3600",
}
return ("", 204, headers)
# Set CORS headers for all responses
headers = {"Access-Control-Allow-Origin": "*"}
def create_error_response(
message: str, status_code: int
) -> tuple[Any, int, dict[str, str]]:
"""Standardize the error responses with common headers."""
return (jsonify({"error": message}), status_code, headers)
# Handle the request and get the search_term
request_json = http_request.get_json(silent=True)
request_args = http_request.args
if request_json and "search_term" in request_json:
search_term = request_json["search_term"]
elif request_args and "search_term" in request_args:
search_term = request_args["search_term"]
else:
return create_error_response("No search term provided", 400)
# Handle the Vertex AI Search and return JSON
try:
results = vertex_ai_search_client.search(search_term)
return (jsonify(results), 200, headers)
except GoogleAPICallError as e:
return create_error_response(
f"Error calling Vertex AI Search API: {str(e)}", 500
)
except ValueError as e:
return create_error_response(f"Invalid input: {str(e)}", 400)
if __name__ == "__main__":
app = Flask(__name__)
@app.route("/", methods=["POST"])
def index() -> tuple[Any, int, dict[str, str]]:
"""
Flask route for handling POST requests when running locally.
This function is used when the script is run directly (not as a Google Cloud Function).
It mimics the behavior of the vertex_ai_search function for local testing.
Returns:
Tuple[Any, int, Dict[str, str]]: The vertex search result.
"""
return vertex_ai_search(request)
app.run("localhost", 8080, debug=True)
search/cloud-function/python/requirements.txt
+7
View File
@@ -0,0 +1,7 @@
Flask==3.1.1
functions_framework==3.8.2
google-cloud-discoveryengine>=0.11
mypy==1.15.0
protobuf==5.29.5
pytest==8.3.5
python-dotenv==1.1.0
search/cloud-function/python/setup_env.sh
Executable
+6
View File
@@ -0,0 +1,6 @@
#!/bin/bash
cp .env.example .env
PROJECT_ID=$(gcloud config get-value project)
sed -i.old "s/PROJECT_ID=.*/PROJECT_ID=$PROJECT_ID/" .env && rm .env.old
echo "Project ID set to $PROJECT_ID in .env file"
echo "Please open .env and set the LOCATION, DATA_STORE_ID, and enum values"
search/cloud-function/python/test_integration_vertex_search_client.py
+121
View File
@@ -0,0 +1,121 @@
# Copyright 2024 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
#
# http://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.
"""
Integration tests for the VertexAISearchClient.
This module contains integration tests that interact with the actual
Vertex AI Search API. These tests require proper configuration of
environment variables and access to the Vertex AI Search service.
"""
from collections.abc import Generator
import os
import pytest
from vertex_ai_search_client import VertexAISearchClient, VertexAISearchConfig
# Load environment variables
PROJECT_ID = os.getenv("PROJECT_ID", "your-project")
LOCATION = os.getenv("LOCATION", "global")
DATA_STORE_ID = os.getenv("DATA_STORE_ID", "your-data-store")
ENGINE_DATA_TYPE = os.getenv("ENGINE_DATA_TYPE", "UNSTRUCTURED")
ENGINE_CHUNK_TYPE = os.getenv("ENGINE_CHUNK_TYPE", "CHUNK")
SUMMARY_TYPE = os.getenv("SUMMARY_TYPE", "VERTEX_AI_SEARCH")
@pytest.fixture(scope="module")
def vertex_ai_search_client() -> Generator[VertexAISearchClient, None, None]:
"""
Fixture to create and yield a VertexAISearchClient instance for testing.
This fixture creates a VertexAISearchClient instance using the
environment variables and yields it for use in tests. The client
is shared across all tests in the module for efficiency.
Yields:
VertexAISearchClient: An instance of the VertexAISearchClient for testing.
"""
config = VertexAISearchConfig(
project_id=PROJECT_ID,
location=LOCATION,
data_store_id=DATA_STORE_ID,
engine_data_type="UNSTRUCTURED",
engine_chunk_type="DOCUMENT_WITH_EXTRACTIVE_SEGMENTS",
summary_type="VERTEX_AI_SEARCH",
)
client = VertexAISearchClient(config)
yield client
def test_search_integration(client: VertexAISearchClient) -> None:
"""
Test the search functionality of VertexAISearchClient with the actual API.
This test performs a search using the VertexAISearchClient and verifies
that the results have the expected structure and content types.
Args:
vertex_ai_search_client (VertexAISearchClient): The client instance to test.
"""
# Perform a search
query = "test query"
results = client.search(query)
# Check the structure of the results
assert "simplified_results" in results
assert isinstance(results["simplified_results"], list)
if results["simplified_results"]:
first_result = results["simplified_results"][0]
assert "metadata" in first_result
assert "page_content" in first_result
# Check for other expected fields
assert "total_size" in results
assert isinstance(results["total_size"], int)
if "summary" in results:
assert "summary_text" in results["summary"]
def test_unstructured_summary() -> None:
"""
Test VertexAISearchClient with unstructured data and summary generation.
This test creates a new VertexAISearchClient instance with specific
settings for unstructured data and summary generation, then performs
a search to verify the results.
"""
config = VertexAISearchConfig(
project_id=PROJECT_ID,
location=LOCATION,
data_store_id=DATA_STORE_ID,
engine_data_type="UNSTRUCTURED",
engine_chunk_type="DOCUMENT_WITH_EXTRACTIVE_SEGMENTS",
summary_type="VERTEX_AI_SEARCH",
)
client = VertexAISearchClient(config)
results = client.search("What is the name of the company?")
# Check the structure of the results
assert "simplified_results" in results
assert isinstance(results["simplified_results"], list)
if results["simplified_results"]:
first_result = results["simplified_results"][0]
assert "metadata" in first_result
assert "page_content" in first_result
if __name__ == "__main__":
pytest.main()
search/cloud-function/python/test_vertex_ai_search_client.py
+338
View File
@@ -0,0 +1,338 @@
# Copyright 2024 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
#
# http://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.
# pylint: disable=redefined-outer-name,protected-access
"""
Unit tests for the VertexAISearchClient class.
This module contains unit tests for the VertexAISearchClient class, using
mocks to simulate the behavior of the Google Cloud Search API. These tests
ensure that the client correctly handles various scenarios and data structures.
"""
import json
from unittest.mock import MagicMock, patch
from google.cloud import discoveryengine_v1alpha as discoveryengine
from google.cloud.discoveryengine_v1alpha.services.search_service.pagers import (
SearchPager,
)
from google.cloud.discoveryengine_v1alpha.types import Document, SearchResponse
import pytest
from vertex_ai_search_client import VertexAISearchClient, VertexAISearchConfig
# Test helper functions
def create_mock_search_pager_result() -> MagicMock:
"""Create a mock SearchPager result for testing."""
mock_pager = MagicMock(spec=SearchPager)
mock_pager.__iter__.return_value = [create_mock_search_pager_return_value()]
mock_pager.total_size = 1
mock_pager.attribution_token = "test-token"
mock_pager.next_page_token = "next-page"
mock_pager.corrected_query = "corrected query"
mock_pager.summary = SearchResponse.Summary(summary_text="Test summary")
mock_pager.applied_controls = None
mock_pager.facets = []
mock_pager.guided_search_result = []
mock_pager.query_expansion_info = []
return mock_pager
def create_mock_search_pager_return_value() -> SearchResponse.SearchResult:
"""Create a mock SearchResponse.SearchResult for testing."""
search_result = SearchResponse.SearchResult()
document = Document()
derived_struct_data = {
"title": "Employee Benefits Summary",
"link": "gs://company-docs/Employee_Benefits_Summary.pdf",
"extractive_answers": [{"content": "Test content", "pageNumber": "1"}],
"snippets": [{"snippet_status": "SUCCESS", "snippet": "Test snippet"}],
}
document.derived_struct_data = derived_struct_data
search_result.document = document
return search_result
# Fixtures
@pytest.fixture
def mock_search_service_client() -> MagicMock:
"""Fixture to create a mock SearchServiceClient."""
with patch(
"vertex_ai_search_client.discoveryengine.SearchServiceClient"
) as mock_client:
mock_client.return_value.serving_config_path.return_value = (
"projects/test-project/locations/us-central1/dataStores/test-data-store/"
"servingConfigs/default_config"
)
yield mock_client
@pytest.fixture
def search_config() -> VertexAISearchConfig:
"""Fixture to create a VertexAISearchConfig instance for testing."""
return VertexAISearchConfig(
project_id="test-project",
location="us-central1",
data_store_id="test-data-store",
engine_data_type="UNSTRUCTURED",
engine_chunk_type="DOCUMENT_WITH_EXTRACTIVE_SEGMENTS",
summary_type="VERTEX_AI_SEARCH",
)
@pytest.fixture
def search_client(search_config: VertexAISearchConfig) -> VertexAISearchClient:
"""Fixture to create a VertexAISearchClient instance for testing."""
with patch(
"vertex_ai_search_client.discoveryengine.SearchServiceClient"
) as mock_client:
mock_client.return_value.serving_config_path.return_value = (
"projects/test-project/locations/us-central1/dataStores/test-data-store/"
"servingConfigs/default_config"
)
return VertexAISearchClient(search_config)
# Tests
def test_init(search_config: VertexAISearchConfig) -> None:
"""Test the initialization of VertexAISearchClient."""
with patch("vertex_ai_search_client.discoveryengine.SearchServiceClient"):
client = VertexAISearchClient(search_config)
assert client.config.project_id == "test-project"
assert client.config.location == "us-central1"
assert client.config.data_store_id == "test-data-store"
assert client.config.engine_data_type == "UNSTRUCTURED"
assert client.config.engine_chunk_type == "DOCUMENT_WITH_EXTRACTIVE_SEGMENTS"
assert client.config.summary_type == "VERTEX_AI_SEARCH"
def test_get_serving_config(search_client: VertexAISearchClient) -> None:
"""Test the get_serving_config method of VertexAISearchClient."""
expected_serving_config = (
"projects/test-project/locations/us-central1/dataStores/test-data-store/"
"servingConfigs/default_config"
)
assert search_client.serving_config == expected_serving_config
def test_build_search_request(search_client: VertexAISearchClient) -> None:
"""Test the build_search_request method of VertexAISearchClient."""
query = "test query"
page_size = 5
# Access to protected member is necessary for testing
request = search_client.build_search_request(
query, page_size
) # pylint: disable=protected-access
assert isinstance(request, discoveryengine.SearchRequest)
assert request.serving_config == search_client.serving_config
assert request.query == query
assert request.page_size == page_size
assert (
request.query_expansion_spec.condition
== discoveryengine.SearchRequest.QueryExpansionSpec.Condition.AUTO
)
assert (
request.spell_correction_spec.mode
== discoveryengine.SearchRequest.SpellCorrectionSpec.Mode.AUTO
)
assert request.content_search_spec.snippet_spec.return_snippet is True
assert request.content_search_spec.summary_spec.summary_result_count == 5
assert request.content_search_spec.summary_spec.include_citations is True
assert request.content_search_spec.summary_spec.ignore_adversarial_query is True
assert (
request.content_search_spec.summary_spec.ignore_non_summary_seeking_query
is True
)
assert (
request.content_search_spec.extractive_content_spec.max_extractive_answer_count
== 1
)
assert (
request.content_search_spec.extractive_content_spec.return_extractive_segment_score
is True
)
def test_map_search_pager_to_dict_basic(search_client: VertexAISearchClient) -> None:
"""Test the map_search_pager_to_dict method with basic data."""
mock_pager = create_mock_search_pager_result()
# Access to protected member is necessary for testing
result = search_client.map_search_pager_to_dict(
mock_pager
) # pylint: disable=protected-access
assert "results" in result
assert len(result["results"]) == 1
assert result["total_size"] == 1
assert result["attribution_token"] == "test-token"
assert result["next_page_token"] == "next-page"
assert result["corrected_query"] == "corrected query"
assert result["summary"]["summary_text"] == "Test summary"
def test_map_search_pager_to_dict_document_content(
search_client: VertexAISearchClient,
) -> None:
"""Test the map_search_pager_to_dict method with document content."""
mock_pager = create_mock_search_pager_result()
# Access to protected member is necessary for testing
result = search_client.map_search_pager_to_dict(
mock_pager
) # pylint: disable=protected-access
document = result["results"][0]["document"]
assert document["derived_struct_data"]["title"] == "Employee Benefits Summary"
assert (
document["derived_struct_data"]["link"]
== "gs://company-docs/Employee_Benefits_Summary.pdf"
)
assert len(document["derived_struct_data"]["extractive_answers"]) == 1
assert (
document["derived_struct_data"]["extractive_answers"][0]["content"]
== "Test content"
)
assert len(document["derived_struct_data"]["snippets"]) == 1
assert document["derived_struct_data"]["snippets"][0]["snippet"] == "Test snippet"
def test_parse_chunk_result(search_client: VertexAISearchClient) -> None:
"""Test the _parse_chunk_result method of VertexAISearchClient."""
chunk = {
"id": "chunk1",
"relevance_score": 0.95,
"content": "Test content",
"document_metadata": {"title": "Test Title", "uri": "https://example.com"},
"page_span": {"page_start": 1, "page_end": 2},
}
# Access to protected member is necessary for testing
result = search_client._parse_chunk_result(
chunk
) # pylint: disable=protected-access
assert result["page_content"] == "Test content"
assert result["metadata"]["chunk_id"] == "chunk1"
assert result["metadata"]["score"] == 0.95
assert result["metadata"]["title"] == "Test Title"
assert result["metadata"]["uri"] == "https://example.com"
assert result["metadata"]["page"] == 1
assert result["metadata"]["page_span_end"] == 2
def test_strip_content() -> None:
"""Test the _strip_content static method of VertexAISearchClient."""
input_text = "<p>Test <strong>content</strong> with &quot;quotes&quot;</p>"
expected_output = 'Test content with "quotes"'
# Access to protected member is necessary for testing
assert (
VertexAISearchClient._strip_content(input_text) == expected_output
) # pylint: disable=protected-access
def test_simplify_search_results_mixed_chunk_and_segments(
search_client: VertexAISearchClient,
) -> None:
"""Test the simplify_search_results method with mixed chunk and segment data."""
input_dict = {
"results": [
{"document": {"id": "doc1", "derived_struct_data": {"title": "Test"}}},
{"chunk": {"id": "chunk1", "content": "Test content"}},
]
}
# Access to protected member is necessary for testing
result = search_client.simplify_search_results(
input_dict
) # pylint: disable=protected-access
assert "simplified_results" in result
assert len(result["simplified_results"]) == 2
assert "metadata" in result["simplified_results"][0]
assert "page_content" in result["simplified_results"][1]
def test_parse_document_result(search_client: VertexAISearchClient) -> None:
"""Test the _parse_document_result method of VertexAISearchClient."""
document = {
"id": "doc1",
"derived_struct_data": {
"title": "Employee Benefits Summary",
"extractive_answers": [
{
"content": "Test content",
"page_number": "3",
}
],
"snippets": [
{
"snippet_status": "SUCCESS",
"snippet": "Test snippet",
}
],
"link": "gs://company-docs/Employee_Benefits_Summary.pdf",
},
}
# Access to protected member is necessary for testing
result = search_client._parse_document_result(
document
) # pylint: disable=protected-access
assert result["metadata"]["title"] == "Employee Benefits Summary"
assert (
result["metadata"]["link"] == "gs://company-docs/Employee_Benefits_Summary.pdf"
)
assert "page_content" in result
assert "Test content" in result["page_content"]
assert "On page 3" in result["page_content"]
@patch("vertex_ai_search_client.VertexAISearchClient.map_search_pager_to_dict")
@patch("vertex_ai_search_client.VertexAISearchClient.simplify_search_results")
def test_search(
mock_simplify: MagicMock,
mock_map_pager: MagicMock,
search_client: VertexAISearchClient,
) -> None:
"""Test the search method of VertexAISearchClient."""
mock_pager = create_mock_search_pager_result()
search_client.client.search.return_value = mock_pager
mock_map_pager.return_value = {"results": [{"document": {"id": "doc1"}}]}
mock_simplify.return_value = {"simplified_results": [{"id": "doc1"}]}
results = search_client.search("test query")
search_client.client.search.assert_called_once()
mock_map_pager.assert_called_once_with(mock_pager)
mock_simplify.assert_called_once_with({"results": [{"document": {"id": "doc1"}}]})
assert results == {"simplified_results": [{"id": "doc1"}]}
results_json = json.dumps(results)
assert results_json == '{"simplified_results": [{"id": "doc1"}]}'
if __name__ == "__main__":
pytest.main()
search/cloud-function/python/vertex_ai_search_client.py
+420
View File
@@ -0,0 +1,420 @@
# Copyright 2024 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
#
# http://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.
"""
VertexAISearchClient for interacting with Google Cloud Vertex AI Search.
This module provides a client class for simplifying interactions with the
Vertex AI Search API. It handles configuration, query construction, and
result parsing.
Example usage:
config = VertexAISearchConfig(
project_id="your-project",
location="global",
data_store_id="your-data-store",
engine_data_type="UNSTRUCTURED",
engine_chunk_type="CHUNK",
summary_type="VERTEX_AI_SEARCH",
)
client = VertexAISearchClient(config)
results = client.search("your search query")
print(results)
"""
from dataclasses import dataclass
import html
import json
import re
from typing import Any, Literal
from google.api_core.client_options import ClientOptions
from google.cloud import discoveryengine_v1alpha as discoveryengine
from google.cloud.discoveryengine_v1alpha.services.search_service.pagers import (
SearchPager,
)
from google.cloud.discoveryengine_v1alpha.types import SearchResponse
# Define types using string literals, similar to enums.
EngineDataTypeStr = Literal["UNSTRUCTURED", "STRUCTURED", "WEBSITE", "BLENDED"]
EngineChunkTypeStr = Literal[
"DOCUMENT_WITH_SNIPPETS", "DOCUMENT_WITH_EXTRACTIVE_SEGMENTS", "CHUNK"
]
SummaryTypeStr = Literal[
"NONE", "VERTEX_AI_SEARCH", "GENERATE_GROUNDED_ANSWERS", "GEMINI"
]
@dataclass
class VertexAISearchConfig:
"""Config for the Vertex AI Search data store."""
project_id: str
location: str
data_store_id: str
engine_data_type: EngineDataTypeStr | str
engine_chunk_type: EngineChunkTypeStr | str
summary_type: SummaryTypeStr | str
def __post_init__(self) -> None:
"""Validate and convert string inputs to appropriate types."""
self.engine_data_type = self._validate_enum(
self.engine_data_type, EngineDataTypeStr, "UNSTRUCTURED"
)
self.engine_chunk_type = self._validate_enum(
self.engine_chunk_type, EngineChunkTypeStr, "CHUNK"
)
self.summary_type = self._validate_enum(
self.summary_type, SummaryTypeStr, "VERTEX_AI_SEARCH"
)
@staticmethod
def _validate_enum(value: str, enum_type: Any, default: str) -> str:
"""Validate and convert string to enum type."""
if value in enum_type.__args__:
return value
print(f"Warning: Invalid value '{value}'. Using default: '{default}'")
return default
def to_dict(self) -> dict[str, str]:
"""Convert the config to a dictionary."""
return {
"project_id": self.project_id,
"location": self.location,
"data_store_id": self.data_store_id,
"engine_data_type": self.engine_data_type,
"engine_chunk_type": self.engine_chunk_type,
"summary_type": self.summary_type,
}
class VertexAISearchClient:
"""
A client for interacting with Google Cloud Vertex AI Search.
This class provides methods to configure the search engine, perform searches,
and parse the results. It supports different types of data stores and search
configurations.
"""
def __init__(self, config: VertexAISearchConfig):
"""
Initialize the VertexAISearchClient.
Args:
config (VertexAISearchConfig): The configuration for the Vertex AI Search client.
"""
self.config = config
self.client = self._create_client()
self.serving_config = self._get_serving_config()
def _create_client(self) -> discoveryengine.SearchServiceClient:
"""
Create and configure the SearchServiceClient.
Returns:
discoveryengine.SearchServiceClient: The configured client.
"""
client_options = None
if self.config.location != "global":
api_endpoint = f"{self.config.location}-discoveryengine.googleapis.com"
client_options = ClientOptions(api_endpoint=api_endpoint)
return discoveryengine.SearchServiceClient(client_options=client_options)
def _get_serving_config(self) -> str:
"""
Get the serving configuration path for the Vertex AI Search data store.
Returns:
str: The serving configuration path.
"""
return self.client.serving_config_path(
project=self.config.project_id,
location=self.config.location,
data_store=self.config.data_store_id,
serving_config="default_config",
)
def search(self, query: str, page_size: int = 10) -> dict[str, Any]:
"""
Perform a search query using Vertex AI Search.
Args:
query (str): The search query.
page_size (int): Number of results to return per page.
Returns:
dict: Parsed and simplified search results.
"""
request = self.build_search_request(query, page_size)
print(f"<request> {request} </request>")
search_pager = self.client.search(request)
response = self.map_search_pager_to_dict(search_pager)
print(f"<response> {response} </response>")
return self.simplify_search_results(response)
def build_search_request(
self, query: str, page_size: int
) -> discoveryengine.SearchRequest:
"""
Build a SearchRequest object based on the client configuration and query.
Args:
query (str): The search query.
page_size (int): Number of results to return per page.
Returns:
discoveryengine.SearchRequest: The configured search request object.
"""
snippet_spec = None
extractive_content_spec = None
if self.config.engine_chunk_type == "DOCUMENT_WITH_SNIPPETS":
snippet_spec = discoveryengine.SearchRequest.ContentSearchSpec.SnippetSpec(
return_snippet=True,
)
if self.config.engine_chunk_type == "DOCUMENT_WITH_EXTRACTIVE_SEGMENTS":
snippet_spec = discoveryengine.SearchRequest.ContentSearchSpec.SnippetSpec(
return_snippet=True,
)
extractive_content_spec = (
discoveryengine.SearchRequest.ContentSearchSpec.ExtractiveContentSpec(
max_extractive_answer_count=1,
return_extractive_segment_score=True,
)
)
summary_spec = None
if self.config.summary_type == "VERTEX_AI_SEARCH":
summary_spec = discoveryengine.SearchRequest.ContentSearchSpec.SummarySpec(
summary_result_count=5,
include_citations=True,
ignore_adversarial_query=True,
ignore_non_summary_seeking_query=True,
)
return discoveryengine.SearchRequest(
serving_config=self.serving_config,
query=query,
page_size=page_size,
content_search_spec=discoveryengine.SearchRequest.ContentSearchSpec(
snippet_spec=snippet_spec,
extractive_content_spec=extractive_content_spec,
summary_spec=summary_spec,
),
query_expansion_spec=discoveryengine.SearchRequest.QueryExpansionSpec(
condition=discoveryengine.SearchRequest.QueryExpansionSpec.Condition.AUTO,
),
spell_correction_spec=discoveryengine.SearchRequest.SpellCorrectionSpec(
mode=discoveryengine.SearchRequest.SpellCorrectionSpec.Mode.AUTO
),
)
def map_search_pager_to_dict(self, pager: SearchPager) -> dict[str, Any]:
"""
Maps a SearchPager to a dictionary structure, iterativly requesting results.
https://cloud.google.com/python/docs/reference/discoveryengine/latest/google.cloud.discoveryengine_v1alpha.services.search_service.pagers.SearchPager
Args:
pager (SearchPager): The pager returned by the search method.
Returns:
Dict[str, Any]: A dictionary containing the search results and metadata.
"""
output: dict[str, Any] = {
"results": [
SearchResponse.SearchResult.to_dict(result) for result in pager
],
"total_size": pager.total_size,
"attribution_token": pager.attribution_token,
"next_page_token": pager.next_page_token,
"corrected_query": pager.corrected_query,
"facets": [],
"applied_controls": [],
}
if pager.summary:
output["summary"] = SearchResponse.Summary.to_dict(pager.summary)
if pager.facets:
output["facets"] = [
SearchResponse.Facet.to_dict(facet) for facet in pager.facets
]
if pager.guided_search_result:
output["guided_search_result"] = SearchResponse.GuidedSearchResult.to_dict(
pager.guided_search_result
)
if pager.query_expansion_info:
output["query_expansion_info"] = SearchResponse.QueryExpansionInfo.to_dict(
pager.query_expansion_info
)
if pager.applied_controls:
output["applied_controls"] = [
control.strip() for control in pager.applied_controls
]
return output
def simplify_search_results(self, response: dict[str, Any]) -> dict[str, Any]:
"""
Simplify the search results by parsing documents and chunks.
Args:
response (Dict[str, Any]): The raw search response.
Returns:
Dict[str, Any]: The simplified search results.
"""
if "results" not in response:
return response
simplified_results = []
for result in response["results"]:
if "document" in result:
simplified_results.append(
self._parse_document_result(result["document"])
)
elif "chunk" in result:
simplified_results.append(self._parse_chunk_result(result["chunk"]))
response["simplified_results"] = simplified_results
return response
def _parse_document_result(self, document: dict[str, Any]) -> dict[str, Any]:
"""
Parse a single document result from the search response.
This supports both structured and unstructured data, and also supports
extractive segments and answers and snippets.
Args:
document (Dict[str, Any]): The document data from the search result.
Returns:
Dict[str, Any]: The parsed document page_content and metadata.
"""
metadata = {
**document.get("derived_struct_data", {}),
**document.get("struct_data", {}),
}
json_data = document.get("json_data", {})
if isinstance(json_data, str):
try:
json_data = json.loads(json_data)
except json.JSONDecodeError:
print(f"Warning: Failed to parse json_data: {json_data}")
json_data = {}
metadata.update(json_data)
result: dict[str, Any] = {"metadata": metadata}
if self.config.engine_data_type == "STRUCTURED":
structured_data = (
json_data if json_data else document.get("struct_data", {})
)
result["page_content"] = json.dumps(structured_data, indent=2)
for key in structured_data.keys():
result["metadata"].pop(key, None)
elif "extractive_answers" in metadata:
result["page_content"] = self._parse_segments(
metadata.get("extractive_answers", [])
)
elif "snippets" in metadata:
result["page_content"] = self._parse_snippets(metadata.get("snippets", []))
else:
result["page_content"] = metadata.get("content", "")
return result
def _parse_segments(self, segments: list[dict[str, Any]]) -> str:
"""
Parse extractive segments from a single document of search results.
Args:
segments (List[Dict[str, Any]]): The list of extractive segments.
Returns:
str: A formatted string containing page number, score and the text of each segment.
"""
parsed_segments = [
{
"content": self._strip_content(segment.get("content", "")),
"page_number": segment.get("page_number") or segment.get("pageNumber"),
"score": segment.get("score"),
}
for segment in segments
]
return "\n\n".join(
f"On page {segment['page_number']} with a relevance score of {segment['score']}:\n"
f"```\n{segment['content']}\n```"
for segment in parsed_segments
)
def _parse_snippets(self, snippets: list[dict[str, Any]]) -> str:
"""
Parse snippets from a single document of search results.
Args:
snippets (List[Dict[str, Any]]): The list of snippets.
Returns:
str: A formatted string containing the successfully parsed snippets.
"""
return "\n\n".join(
self._strip_content(snippet.get("snippet", ""))
for snippet in snippets
if snippet.get("snippetStatus") == "SUCCESS"
)
def _parse_chunk_result(self, chunk: dict[str, Any]) -> dict[str, Any]:
"""
Parse a single chunk result from the search response.
Args:
chunk (Dict[str, Any]): The chunk data from the search result.
Returns:
Dict[str, Any]: The parsed chunk page_content and metadata.
"""
metadata = {
"chunk_id": chunk.get("id"),
"score": chunk.get("relevance_score"),
}
page_span = chunk.get("page_span", {})
if page_span:
metadata["page"] = page_span.get("page_start")
metadata["page_span_end"] = page_span.get("page_end")
metadata.update(chunk.get("document_metadata", {}))
metadata.update(chunk.get("derived_struct_data", {}))
return {
"metadata": metadata,
"page_content": self._strip_content(chunk.get("content", "")),
}
@staticmethod
def _strip_content(text: str) -> str:
"""
Strip HTML tags and unescape HTML entities from the given text.
Args:
text (str): The input text to clean.
Returns:
str: The cleaned text.
"""
text = re.sub("<.*?>", "", text)
return html.unescape(text).strip()