You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

40 lines
1.1 KiB

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

import json
import time
from kafka import KafkaConsumer
from AiService.Model.TaskModel import TaskModel
from AiService.Config.Config import KAFKA_BOOTSTRAP_SERVERS, KAFKA_TOPIC
def process_task(task_data: dict):
"""
模拟任务处理
"""
task_id = task_data.get("task_id")
keyword = task_data.get("prompt")
# 模拟处理任务(休息 3 秒)
print(f"Processing task {task_id} with keyword: {keyword}")
time.sleep(3)
# 更新任务状态为“已完成”
tm = TaskModel()
tm.update_task_status(task_id, "completed", result_url="http://example.com/result")
tm.close()
print(f"Task {task_id} completed.")
def start_worker():
"""
启动 Worker从 Kafka 中消费任务
"""
consumer = KafkaConsumer(
KAFKA_TOPIC,
bootstrap_servers=KAFKA_BOOTSTRAP_SERVERS,
value_deserializer=lambda v: json.loads(v.decode('utf-8'))
)
print("Worker started. Waiting for tasks...")
for message in consumer:
task_data = message.value
process_task(task_data)
if __name__ == "__main__":
start_worker()