chore: import upstream snapshot with attribution

2026-07-13 13:30:30 +08:00
commit 914fea506e
2793 changed files with 802106 additions and 0 deletions
@@ -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)