wehub-resource-sync
136 lines
4.8 KiB
Python
136 lines
4.8 KiB
Python
# 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)
|