feat: connect to vertex ai

public-prediction/kafka_consumer.py

@@ -4,13 +4,13 @@ from kafka import KafkaConsumer
 from get_gpt_answer import GetGPTAnswer
 from typing import List
 from concurrent.futures import ThreadPoolExecutor
+from predict_custom_model import predict_custom_trained_model
+from google.protobuf.json_format import MessageToDict
 
 
 def get_gpt_responses(data: dict[str, any], gpt_helper: GetGPTAnswer):
-    # data["gpt35_answer"] = gpt_helper.generate_gpt35_answer(data["question"])
-    # data["gpt4_answer"] = gpt_helper.generate_gpt4_answer(data["question"])
-    data["gpt35_answer"] = "This is gpt35 answer"
-    data["gpt4_answer"] = "This is gpt4 answer"
+    data["gpt35_answer"] = gpt_helper.generate_gpt35_answer(data["question"])
+    data["gpt4_answer"] = gpt_helper.generate_gpt4_answer(data["question"])
     return data
 
 
@@ -23,6 +23,22 @@ def process_batch(batch: List[dict[str, any]], batch_size: int, job_application_
 
     print("Batch ready with gpt responses", results)
 
+    predictions = predict_custom_trained_model(
+        instances=results, project=os.environ.get("PROJECT_ID"), endpoint_id=os.environ.get("ENDPOINT_ID"))
+
+    results = []
+    for prediction in predictions:
+        result_dict = {}
+        for key, value in prediction._pb.items():
+            # Ensure that 'value' is a protobuf message
+            if hasattr(value, 'DESCRIPTOR'):
+                result_dict[key] = MessageToDict(value)
+            else:
+                print(f"Item {key} is not a convertible protobuf message.")
+        results.append(result_dict)
+
+    print({"result": results})
+
 
 def consume_messages():
     consumer = KafkaConsumer(

public-prediction/predict_custom_model.py

@@ -0,0 +1,40 @@
+from typing import Dict, List, Union
+from google.cloud import aiplatform
+from google.protobuf import json_format
+from google.protobuf.struct_pb2 import Value
+
+
+def predict_custom_trained_model(
+    project: str,
+    endpoint_id: str,
+    instances: Union[Dict, List[Dict]],
+    location: str = "us-central1",
+    api_endpoint: str = "us-central1-aiplatform.googleapis.com",
+):
+    """
+    `instances` can be either single instance of type dict or a list
+    of instances.
+    """
+    # The AI Platform services require regional API endpoints.
+    client_options = {"api_endpoint": api_endpoint}
+    # Initialize client that will be used to create and send requests.
+    # This client only needs to be created once, and can be reused for multiple requests.
+    client = aiplatform.gapic.PredictionServiceClient(
+        client_options=client_options)
+    # The format of each instance should conform to the deployed model's prediction input schema.
+    instances = instances if isinstance(instances, list) else [instances]
+    instances = [
+        json_format.ParseDict(instance_dict, Value()) for instance_dict in instances
+    ]
+    parameters_dict = {}
+    parameters = json_format.ParseDict(parameters_dict, Value())
+    endpoint = client.endpoint_path(
+        project=project, location=location, endpoint=endpoint_id
+    )
+    response = client.predict(
+        endpoint=endpoint, instances=instances, parameters=parameters
+    )
+    # The predictions are a google.protobuf.Value representation of the model's predictions.
+    predictions = response.predictions
+
+    return predictions