wehub-resource-sync
152 lines
5.2 KiB
Python
152 lines
5.2 KiB
Python
# Copyright 2025 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.
|
|
|
|
import json
|
|
from datetime import datetime
|
|
from threading import Thread
|
|
from uuid import uuid4
|
|
|
|
from flask import Request, jsonify
|
|
from google.cloud.aiplatform import MatchingEngineIndex, MatchingEngineIndexEndpoint
|
|
from langchain_google_vertexai import VertexAIEmbeddings
|
|
from src.bigquery import (
|
|
EMBEDDINGS_TABLE,
|
|
INTENTS_TABLE,
|
|
INTENTS_TABLE_ID_COLUMN,
|
|
BigQueryRepository,
|
|
)
|
|
from src.chunk import ChunkService
|
|
from src.cloud_storage import EMBEDDINGS_FILE, EMBEDDINGS_FOLDER, CloudStorageRepository
|
|
from src.models import Embedding, Intent
|
|
|
|
INDEX_DIMENSIONS = 768
|
|
INDEX_DISTANCE_MEASURE = "DOT_PRODUCT_DISTANCE"
|
|
INDEX_NEIGHBORS_COUNT = 150
|
|
|
|
TEXT_EMBEDDING_MODEL = "textembedding-gecko@003"
|
|
EMBEDDINGS_MODEL = VertexAIEmbeddings(TEXT_EMBEDDING_MODEL)
|
|
|
|
TIME_FORMAT = "%Y-%m-%d %H:%M:%S"
|
|
|
|
|
|
def create_intent_index(request: Request):
|
|
if request.method != "POST":
|
|
return jsonify({"error": "Method not allowed"}), 405
|
|
try:
|
|
request_json = request.get_json()
|
|
intent_name = request_json.get("intent_name")
|
|
index_resource = request_json.get("index_endpoint_resource")
|
|
except Exception:
|
|
return jsonify({"error": "Bad Request"}), 400
|
|
|
|
print(f"Event decoded {request_json}", intent_name, index_resource)
|
|
big_query_repository = BigQueryRepository()
|
|
gcs_repository = CloudStorageRepository(big_query_repository.client.project)
|
|
|
|
try:
|
|
results = big_query_repository.get_row_by_id(
|
|
INTENTS_TABLE, INTENTS_TABLE_ID_COLUMN, intent_name
|
|
)
|
|
intent = None
|
|
for row in results:
|
|
intent = Intent.__from_row__(row)
|
|
|
|
index_endpoint = MatchingEngineIndexEndpoint(index_resource)
|
|
|
|
print("Everything has been correctly received")
|
|
|
|
index_embeddings = ""
|
|
chunk_service = ChunkService(
|
|
big_query_repository.client.project, intent.gcp_bucket
|
|
)
|
|
embeddings = []
|
|
|
|
index_unique_name = (
|
|
f"{intent.name.lower().replace(' ', '-').replace('_', '-')}-{uuid4()}"
|
|
)
|
|
chunks = chunk_service.generate_chunks()
|
|
|
|
for index, chunk in enumerate(chunks):
|
|
embedding = create_embeddings(chunk)
|
|
|
|
if embedding is not None:
|
|
doc_id = f"{intent.name}-{index}.txt"
|
|
embeddings.append(
|
|
Embedding(
|
|
id=doc_id,
|
|
text=chunk,
|
|
index=index_unique_name,
|
|
author="system",
|
|
timestamp=datetime.now().strftime(TIME_FORMAT),
|
|
)
|
|
)
|
|
index_embeddings += (
|
|
json.dumps(
|
|
{
|
|
"id": doc_id,
|
|
"embedding": [str(value) for value in embedding],
|
|
}
|
|
)
|
|
+ "\n"
|
|
)
|
|
print(f"Embeddings created for {[e.id for e in embeddings]}")
|
|
print(f"Uploading embeddings {intent.name}/{EMBEDDINGS_FILE}")
|
|
gcs_repository.create(
|
|
f"{EMBEDDINGS_FOLDER}/{intent.name}/{EMBEDDINGS_FILE}", index_embeddings
|
|
)
|
|
|
|
index = create_index(
|
|
index_unique_name,
|
|
intent.name,
|
|
gcs_repository.bucket_name,
|
|
)
|
|
big_query_repository.update_intent_status(intent_name, "3")
|
|
print("Uploading text chunks to bigquery...")
|
|
big_query_repository.insert_rows(EMBEDDINGS_TABLE, embeddings)
|
|
Thread(target=deploy_index_endpoint, args=(index_endpoint, index)).start()
|
|
return jsonify({"message": "JSON received and processed"}), 200
|
|
|
|
except Exception as e:
|
|
if index:
|
|
big_query_repository.update_intent_status(intent_name, "4")
|
|
else:
|
|
big_query_repository.update_intent_status(intent_name, "2")
|
|
print(str(e))
|
|
return jsonify({"error": str(e)}), 500
|
|
|
|
|
|
def create_embeddings(chunk: str) -> list[float]:
|
|
return EMBEDDINGS_MODEL.embed_query(chunk)
|
|
|
|
|
|
def create_index(index_unique_name: str, intent_name: str, bucket_name: str):
|
|
print(f"Creating index: {index_unique_name}")
|
|
return MatchingEngineIndex.create_tree_ah_index(
|
|
display_name=index_unique_name,
|
|
dimensions=INDEX_DIMENSIONS,
|
|
approximate_neighbors_count=INDEX_NEIGHBORS_COUNT,
|
|
distance_measure_type=INDEX_DISTANCE_MEASURE,
|
|
contents_delta_uri=f"gs://{bucket_name}/{EMBEDDINGS_FOLDER}/{intent_name}",
|
|
)
|
|
|
|
|
|
def deploy_index_endpoint(
|
|
index_endpoint: MatchingEngineIndexEndpoint, index: MatchingEngineIndex
|
|
):
|
|
print("Deploying index...")
|
|
index_endpoint.deploy_index(
|
|
index=index,
|
|
deployed_index_id=index.display_name.replace("-", "_"),
|
|
)
|