Spaces:
Running
Running
| import json | |
| import os | |
| import requests | |
| 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["gpt4o_answer"] = gpt_helper.generate_gpt4o_answer(data["question"]) | |
| return data | |
| def process_batch(batch: List[dict[str, any]], batch_size: int, gpt_helper: GetGPTAnswer): | |
| with ThreadPoolExecutor(max_workers=batch_size) as executor: | |
| futures = [executor.submit( | |
| get_gpt_responses, data, gpt_helper) for data in batch] | |
| results = [future.result() for future in futures] | |
| 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) | |
| return results | |
| def send_results_back(full_results: dict[str, any], job_application_id: str): | |
| print(f"Sending results back with job_app_id {job_application_id}") | |
| url = "https://ta-2-sistem-cerdas-be-vi2jkj4riq-et.a.run.app/api/anti-cheat/result" | |
| headers = { | |
| "Content-Type": "application/json", | |
| "x-api-key": os.environ.get("X-API-KEY") | |
| } | |
| body = { | |
| "job_application_id": job_application_id, | |
| "evaluations": full_results | |
| } | |
| response = requests.patch(url, json=body, headers=headers) | |
| print(f"Data sent with status code {response.status_code}") | |
| print(response.content) | |
| def consume_messages(): | |
| consumer = KafkaConsumer( | |
| "ai-detector", | |
| bootstrap_servers=[os.environ.get("KAFKA_IP")], | |
| auto_offset_reset='earliest', | |
| client_id="ai-detector-1", | |
| group_id="ai-detector", | |
| api_version=(0, 10, 2) | |
| ) | |
| print("Successfully connected to Kafka at", os.environ.get("KAFKA_IP")) | |
| BATCH_SIZE = 5 | |
| gpt_helper = GetGPTAnswer() | |
| for message in consumer: | |
| try: | |
| incoming_message = json.loads(message.value.decode("utf-8")) | |
| full_batch = incoming_message["data"] | |
| except json.JSONDecodeError: | |
| print("Failed to decode JSON from message:", message.value) | |
| print("Continuing...") | |
| continue | |
| print("Parsing successful. Processing job_app_id {0}".format( | |
| incoming_message['job_application_id'])) | |
| full_results = [] | |
| for i in range(0, len(full_batch), BATCH_SIZE): | |
| batch = full_batch[i:i+BATCH_SIZE] | |
| batch_results = process_batch(batch, BATCH_SIZE, gpt_helper) | |
| full_results.extend(batch_results) | |
| send_results_back(full_results, incoming_message["job_application_id"]) | |