mint/ai-translate/ai_translate/service.py

import json
import time
import logging
from typing import List, Dict, Any, Optional
from datetime import datetime
from dataclasses import dataclass
import time

import requests

from .utils import is_stopped

logger = logging.getLogger(__name__)


class DatabaseException(Exception):
    def __init__(self, message="database api access exception"):
        self.message = message
        super().__init__(self.message)


class SectionTimeout(Exception):
    def __init__(self, message="片段超时"):
        self.message = message
        super().__init__(self.message)


class TaskFailException(Exception):
    def __init__(self, message="task fail"):
        self.message = message
        super().__init__(self.message)


class LLMFailException(Exception):
    def __init__(self, message="LLM request fail"):
        self.message = message
        super().__init__(self.message)


@dataclass
class TaskProgress:
    """任务进度"""
    current: int
    total: int


@dataclass
class TaskInfo:
    """任务模型"""
    id: str
    title: str
    category: str
    description: str
    status: str


@dataclass
class Task:
    """任务模型"""
    info: TaskInfo
    progress: TaskProgress


@dataclass
class AiModel:
    """AI模型配置"""
    uid: str
    model: str
    url: str
    key: str
    token: str
    system_prompt: Optional[str] = None


@dataclass
class Sentence:
    """句子模型"""
    book_id: str
    paragraph: int
    word_start: int
    word_end: int
    channel_uid: str
    content: str
    content_type: str = 'markdown'
    access_token: Optional[str] = None


@dataclass
class Payload:
    """消息模型"""
    model: AiModel
    task: Task
    prompt: str
    sentence: Sentence


@dataclass
class Message:
    """消息模型"""
    model: AiModel
    task: Task
    payload: List[Payload]
    sentence: Sentence


class AiTranslateService:
    """AI翻译服务"""

    def __init__(self, redis, ch, method, api_url, openai_proxy, customer_timeout, worker_name: str):
        self.queue = 'ai_translate'
        self.model_token = None
        self.task = None
        self.redis_clusters = redis[0]
        self.redis_namespace = redis[1]
        self.api_timeout = 100
        self.llm_timeout = 300
        self.task_topic_id = None
        self.api_url = api_url
        self.customer_timeout = customer_timeout
        self.channel = ch
        self.maxProcessTime = 15 * 60  # 一个句子的最大处理时间
        self.openai_proxy = openai_proxy
        self.worker_name = worker_name

    def process_translate(self, message_id: str, body: Message) -> bool:
        """处理翻译任务"""
        is_stopped()
        taskStartAt = int(time.time())

        self.task = body.task

        self.redis_clusters.set(
            f"{self.redis_namespace}/task/{self.task.id}/message_id", message_id)
        pointer_key = f"{self.redis_namespace}/task/{self.task.id}/pointer"
        pointer = 0

        if self.redis_clusters.exists(pointer_key):
            # 回到上次中断的点
            pointer = int(self.redis_clusters.get(pointer_key))
            logger.info(f"last break point {pointer}")
            if pointer >= len(body.payload):
                self.redis_clusters.delete(pointer_key)
                return True

        # 获取model token
        self.model_token = body.model.token

        self._set_task_status(self.task.id, 'running')

        # 设置task discussion topic
        taskTopicKey = f'{self.redis_namespace}/message/{message_id}/topic'
        if self.redis_clusters.exists(taskTopicKey):
            # 获取上次的task topic id
            self.task_topic_id = self.redis_clusters.get(taskTopicKey)
        else:
            self.task_topic_id = self._task_discussion(
                self.task.id,
                'task',
                self.task.title,
                f'id:{message_id} worker:{self.worker_name}',
                None
            )
        times = [self.maxProcessTime]
        # breakpoint()
        for i in range(pointer, len(body.payload)):
            is_stopped()

            startAt = int(time.time())

            message = body.payload[i]
            task_discussion_content = []

            # 推理

            response_llm = self._request_llm(message)
            task_discussion_content.append('- LLM request successful')

            if self.task.category == 'translate':
                # 写入句子库
                message.sentence.content = response_llm['content']
                self._save_sentence(message.sentence)

            if self.task.category == 'suggest':
                # 写入pr
                self._save_pr(message.sentence, response_llm['content'])

            # 获取句子id
            s_uid = self._get_sentence_id(message.sentence)

            # 写入句子 discussion
            if s_uid:
                topic_children = []
                # 任务结果
                topic_children.append(response_llm['content'])
                # 推理过程写入discussion
                if response_llm.get('reasoningContent'):
                    topic_children.append(response_llm['reasoningContent'])
                self._sentence_discussion(
                    s_uid, message.prompt, topic_children)

            # 修改task 完成度
            progress = self._set_task_progress(
                TaskProgress(i + 1, len(body.payload)))
            task_discussion_content.append(f"- progress={progress}")

            # 写入task discussion
            if self.task_topic_id:
                content = '\n'.join(task_discussion_content)
                d_id = self._task_discussion(
                    self.task.id,
                    'task',
                    self.task.title,
                    content,
                    self.task_topic_id
                )
            else:
                logger.error('no task discussion root')

        # 任务完成 修改任务状态为 done
        self._set_task_status(self.task.id, 'done')
        self.redis_clusters.delete(pointer_key)
        logger.info('ai translate task complete')
        return True

    def _sentence_discussion(self, id, prompt, discussions):
        topic_id = self._task_discussion(
            id,
            'sentence',
            self.task.title,
            prompt,
            None
        )

        if topic_id:
            logger.info(f'{self.queue} discussion create topic successful')

            for content in discussions:
                logger.debug(f'{self.queue} discussion child request')
                d_id = self._task_discussion(
                    id, 'sentence', self.task.title, content, topic_id)
                if d_id:
                    logger.info(
                        f'{self.queue} discussion child successful')
        else:
            logger.error(
                f'{self.queue} discussion create topic response is null')

    def _set_task_status(self, task_id: str, status: str):
        """设置任务状态"""
        url = f"{self.api_url}/v2/task-status/{task_id}"
        data = {'status': status}
        logger.debug(f'ai_translate task status request: {url}, data: {data}')

        headers = {'Authorization': f'Bearer {self.model_token}'}
        response = requests.patch(
            url, json=data, headers=headers, timeout=self.api_timeout)

        if response.ok:
            logger.info(f'ai_translate task status successful ({status})')
        else:
            logger.error(
                f'ai_translate task status update fail. response: {response.text}')

    def _save_model_log(self, token: str, data: Dict[str, Any]) -> bool:
        """保存模型日志"""
        url = f"{self.api_url}/v2/model-log"

        headers = {'Authorization': f'Bearer {token}'}
        response = requests.post(
            url, json=data, headers=headers, timeout=self.api_timeout)
        if not response.ok:
            logger.error(
                f'ai-translate model log create failed: {response.text}')
            return False
        return True

    def _task_discussion(self, res_id: str, res_type: str, title: str, content: str, parent_id: Optional[str] = None):
        """创建任务讨论"""
        url = f"{self.api_url}/v2/discussion"

        task_discussion_data = {
            'res_id': res_id,
            'res_type': res_type,
            'content': content,
            'content_type': 'markdown',
            'type': 'discussion',
            'notification': False,
        }

        if parent_id:
            task_discussion_data['parent'] = parent_id
        else:
            task_discussion_data['title'] = title

        logger.info(f'{self.queue} discussion create: {url},')

        headers = {'Authorization': f'Bearer {self.model_token}'}
        response = requests.post(
            url, json=task_discussion_data, headers=headers, timeout=self.api_timeout)

        if not response.ok:
            logger.error(
                f'{self.queue} discussion create error: {response.json()}')
            return False

        # logger.debug(
        #    f'{self.queue} discussion create: {json.dumps(response.json())}')

        response_data = response.json()
        if response_data.get('data', {}).get('id'):
            return response_data['data']['id']
        return False

    def _request_llm(self, message: Message) -> Dict[str, Any]:
        """请求LLM"""
        param = {
            "model": message.model.model,
            "messages": [
                {"role": "system", "content": message.model.system_prompt or ''},
                {"role": "user", "content": message.prompt},
            ],
            "temperature": 0.7,
            "stream": False
        }

        logger.info(
            f'{self.queue} LLM request {message.model.url} model: {param["model"]}')

        # logger.debug(
        #     f'{self.queue} LLM api request: {message.model.url}, data: {json.dumps(param)}')

        # 写入 model log
        model_log_data = {
            'model_id': message.model.uid,
            'request_at': datetime.now().isoformat(),
            'request_data': json.dumps(param, ensure_ascii=False),
        }

        # 失败重试
        max_retries = 3
        attempt = 0

        headers = {'Authorization': f'Bearer {message.model.key}'}

        while attempt < max_retries:
            try:
                if self.openai_proxy:
                    response = requests.post(
                        self.openai_proxy,
                        json={
                            "open_ai_url": message.model.url,
                            "api_key": message.model.key,
                            'payload': param,
                        },
                        headers=headers,
                        timeout=self.llm_timeout
                    )
                else:
                    response = requests.post(
                        message.model.url,
                        json=param,
                        headers=headers,
                        timeout=self.llm_timeout
                    )
                response.raise_for_status()

                logger.info(f'{self.queue} LLM request successful')

                model_log_data.update({
                    'request_headers': json.dumps(dict(response.request.headers), ensure_ascii=False),
                    'response_headers': json.dumps(dict(response.headers), ensure_ascii=False),
                    'status': response.status_code,
                    'response_data': json.dumps(response.json(), ensure_ascii=False),
                    'success': True
                })
                break
            except requests.exceptions.RequestException as e:
                model_log_data.update({
                    'request_headers': json.dumps(dict(e.response.request.headers), ensure_ascii=False),
                    'response_headers': json.dumps(dict(e.response.headers), ensure_ascii=False),
                    'status': e.response.status_code,
                    'response_data': e.response.text,
                    'success': False
                })
                attempt += 1
                status = getattr(e.response, 'status_code',
                                 0) if hasattr(e, 'response') else 0

                # 某些错误不需要重试
                if status in [400, 401, 403, 404, 422]:
                    logger.warning(f"客户端错误，不重试: {status}")
                    self._save_model_log(self.model_token, model_log_data)
                    raise LLMFailException

                # 服务器错误或网络错误可以重试
                if attempt < max_retries:
                    delay = 2 ** attempt  # 指数退避
                    logger.warning(f"请求失败（第 {attempt} 次），{delay} 秒后重试...")
                    time.sleep(delay)
                else:
                    logger.error("达到最大重试次数，请求最终失败")
                    raise e
            except Exception as e:
                raise e
            finally:
                try:
                    self._save_model_log(self.model_token, model_log_data)
                    logger.info(f'{self.queue} model log saved')
                except Exception as e:
                    logger.error(e)

        ai_data = response.json()

        response_content = ai_data['choices'][0]['message']['content']
        reasoning_content = ai_data['choices'][0]['message'].get(
            'reasoning_content')

        output = {'content': response_content}
        logger.debug(f'{self.queue} LLM response content={response_content}')

        if not reasoning_content:
            logger.debug(f'{self.queue} no reasoningContent')
        else:
            logger.debug(f'{self.queue} reasoning={reasoning_content}')
            output['reasoningContent'] = reasoning_content

        return output

    def _save_sentence(self, sentence: Sentence):
        """写入句子库"""
        url = f"{self.api_url}/v2/sentence"

        logger.info(f"{self.queue} sentence update {url}")

        headers = {'Authorization': f'Bearer {self.model_token}'}
        data = {'sentences': [sentence.__dict__]}

        response = requests.post(
            url, json=data, headers=headers, timeout=self.api_timeout)

        if not response.ok:
            logger.error(
                f'{self.queue} sentence update failed: {url}, data: {response.json()}')
            raise DatabaseException("sentence 数据库写入错误")

        count = response.json()['data']['count']
        logger.info(f"{self.queue} sentence update {count} successful")

    def _save_pr(self, sentence: Sentence, content: str):
        """保存PR"""
        url = f"{self.api_url}/v2/sentpr"
        logger.info(f"{self.queue} sentence update {url}")

        data = {
            'book': sentence.book_id,
            'para': sentence.paragraph,
            'begin': sentence.word_start,
            'end': sentence.word_end,
            'channel': sentence.channel_uid,
            'text': content,
            'notification': False,
            'webhook': False,
        }

        headers = {'Authorization': f'Bearer {self.model_token}'}
        response = requests.post(
            url, json=data, headers=headers, timeout=self.api_timeout)

        if not response.ok:
            logger.error(
                f'{self.queue} sentence update failed: {url}, data: {response.json()}')
            raise DatabaseException("pr 数据库写入错误")

        if response.json().get('ok'):
            logger.info(f"{self.queue} sentence suggest update successful")
        else:
            logger.error(
                f"{self.queue} sentence suggest update failed: {url}, data: {response.json()}")

    def _get_sentence_id(self, sentence: Sentence) -> str:
        """获取句子ID"""
        try:
            url = f"{self.api_url}/v2/sentence-info/aa"
            logger.info(f'ai translate: {url}')

            params = {
                'book': sentence.book_id,
                'par': sentence.paragraph,
                'start': sentence.word_start,
                'end': sentence.word_end,
                'channel': sentence.channel_uid
            }

            headers = {'Authorization': f'Bearer {self.model_token}'}
            response = requests.get(
                url, params=params, headers=headers, timeout=self.api_timeout)

            if not response.json().get('ok'):
                logger.error(
                    f'{self.queue} sentence id error: {response.text}')
                return False

            s_uid = response.json()['data']['id']
            logger.debug(f"sentence id={s_uid}")
            return s_uid
        except Exception as e:
            logger.error(f"error: {e}")
            return False

    def _set_task_progress(self, current: TaskProgress) -> int:
        """设置任务进度"""
        if current.total > 0:
            progress = int(current.current * 100 / current.total)
        else:
            progress = 100
            logger.error(
                f'{self.queue} progress total is zero, task_id: {self.task.id}')

        url = f"{self.api_url}/v2/task/{self.task.id}"
        data = {'progress': progress}

        logger.debug(
            f'{self.queue} task progress request: {url}, data: {data}')

        headers = {'Authorization': f'Bearer {self.model_token}'}
        response = requests.patch(
            url, json=data, headers=headers, timeout=self.api_timeout)

        if not response.ok:
            logger.error(
                f'{self.queue} task progress error: {response.json()}')
        else:
            logger.info(
                f'{self.queue} task progress successful progress={response.json()["data"]["progress"]}')

        return progress

    def handle_failed(self, message_id: str, message: str, exception: Exception):
        """处理失败的翻译任务"""
        try:
            # 彻底失败时的业务逻辑
            # 设置task为失败状态
            self._set_task_status(self.task.id, 'stop')

            # 将故障信息写入task discussion
            if self.task_topic_id:
                error_message = f"**任务处理失败** 请重启任务 \n- message id={message_id} \n- 错误信息：{message} \n- 异常：{str(exception)}"
                d_id = self._task_discussion(
                    self.task.id,
                    'task',
                    '任务处理失败',
                    error_message,
                    self.task_topic_id
                )
        except Exception as e:
            logger.error(f'处理失败ai任务时出错: {str(e)}')

    def handle_retry(self, message_id: str, message: str, exception: Exception):
        """处理失败 需要重试"""
        try:
            # 失败时的业务逻辑
            self._set_task_status(self.task.id, 'pause')
            # 将故障信息写入task discussion
            if self.task_topic_id:
                error_message = f"任务处理出错 正在重试 \n- message id={message_id} \n- 错误信息：{message} \n- 异常：{str(exception)}"
                d_id = self._task_discussion(
                    self.task.id,
                    'task',
                    '任务处理出错',
                    error_message,
                    self.task_topic_id
                )
        except Exception as e:
            logger.error(f'处理失败ai任务时出错: {str(e)}')

    def handle_complete(self):
        try:
            # 将故障信息写入task discussion
            if self.task_topic_id:
                d_id = self._task_discussion(
                    self.task.id,
                    'task',
                    '任务处理完成',
                    '任务处理完成',
                    self.task_topic_id
                )
        except Exception as e:
            logger.error(f'处理任务完成时出错: {str(e)}')

    def get_task_id(self) -> str:
        return self.task.id