QingLong/AI/AiService/Worker/TaskWorker.py

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()