src/agent.py

import os
import time

import gradio as gr
from phi.assistant import Assistant
from phi.llm.openai import OpenAIChat
from pydantic import BaseModel

from .log import logger

qwen_model = 'qwen-plus'
qwen_api_key = os.getenv("QWEN_DS_API_KEY")
qwen_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1'


class TranslateFormat(BaseModel):
    en: str


class PormptExtendFormat(BaseModel):
    prompt: str


trans_assist = Assistant(
    llm=OpenAIChat(model=qwen_model, max_tokens=3000, temperature=0.3, api_key=qwen_api_key, base_url=qwen_url),
    description='你是专业的英语翻译，请将用户输入的一句中文翻译成英文',
    output_model=TranslateFormat
)

requirements = [
    "1.对于过于简短的用户输入，在不改变原意及主体外观的前提下，合理补充具象化细节，避免使用过于抽象描述。",
    "2.如果涉及人像，完善与用户输入相符的人物外貌、表情、种族、姿态、景别、穿着、影调、质感等方面描述。",
    "3.匹配符合用户意图且精准详细的风格描述。如果用户未指定，则根据画面选择最恰当的风格，或使用纪实摄影照片。",
    "4.若输入为中文，则整体中文输出；若输入为全英文，则整体英文输出；保留引号中原文以及重要的输入信息，不要改写。",
    "5.在语义完整前提下，改写后prompt字数小于200字，避免prompt冗长。",
    """
    6.改写后 prompt 示例：
    (1) 日系小清新写真照片，扎着双麻花辫的波西米亚小女孩坐在船边。女孩穿着白色方领泡泡袖连衣裙，裙子上有褶皱和纽扣装饰。她皮肤白皙，五官精致，眼泪汪汪直视镜头。她双手扶船，长发刘海遮住部分额头。背景是清澈明亮的户外场景，可见蓝天、山峦和一些干枯植物。高清写实摄影，近景中心对称构图。\n
    (2) 二次元厚涂动漫插画，一个猫耳东亚萌妹手持文件夹，怒气冲冲走向电脑。她深紫色爆炸头，红色眼睛，头顶有一个粉色光圈。少女身穿深灰色短裙和浅灰色上衣，腰间系着白色系带，胸前佩戴名牌，居中写着黑体中文"紫阳"。室内背景中摆放了很多办公桌。粗线条的日系赛璐璐风格。近景半身略仰视视角。\n
    (3) 美剧艺术海报风格，身穿黄色防护服的老年Walter White坐在金属折叠椅上，头顶无衬线英文写着"Breaking Bad"，周围是成堆的美元和蓝色塑料储物箱。他戴着眼镜目光直视镜头，左手拿着一支雪茄，右手放在膝盖上。背景是废弃阴暗的厂房，阳光透过窗户照射进来。画面带有明显颗粒质感纹理。长焦人物平视特写。
    (4) CG game concept digital art featuring three giant mutant crocodiles with wide-open mouths, revealing pink tongues and sharp teeth. Their rough skin resembles grayish-white stone. On the back of the crocodile to the left, lush trees, shrubs, and some thorn-like protrusions grow. The background of the scene shows a dusk sky and a shimmering pond. The overall atmosphere is dark and cold, with a tilted composition that creates a strong sense of depth and layers. \n
    """
]

prompt_writer = Assistant(
    llm=OpenAIChat(model=qwen_model, max_tokens=8000, temperature=0.3, api_key=qwen_api_key, base_url=qwen_url),
    description="你是一位prompt优化师，旨在将用户输入改写为优质prompt，在不影响原意前提下更高质量生图。",
    instructions=requirements,
    debug_mode=False,
    output_model=PormptExtendFormat,
)


def contains_chinese(text):
    """
    判断文本中是否包含中文字符（基于 Unicode 范围）
    """
    # 定义中文的 Unicode 范围（基础汉字 + 扩展区）
    cjk_ranges = [
        (0x4e00, 0x9fff),  # CJK 基本汉字区（含常用汉字）
        (0x3400, 0x4dbf),  # CJK 扩展 A
        (0x20000, 0x2a6df),  # CJK 扩展 B
        (0x2a700, 0x2b73f),  # CJK 扩展 C
        (0x2b740, 0x2b81f),  # CJK 扩展 D
        (0x2b820, 0x2ceaf),  # CJK 扩展 E
    ]

    for c in text:
        code = ord(c)
        for start, end in cjk_ranges:
            if start <= code <= end:
                return True
    return False


def translate_prompt(prompt, max_attempts=5):
    attempts = 0
    while attempts < max_attempts:
        try:
            if contains_chinese(prompt):
                res = trans_assist.run(prompt)
                logger.info(f"translate Chinese prompt into English: {prompt} -> {res.en}")
                return res.en
            else:
                return prompt
        except Exception as e:
            print(f"尝试 {attempts + 1} 失败：{e}")
            attempts += 1
            if attempts < 5:
                time.sleep(1)  # 延迟重试
    raise gr.Error(f"提示词扩写达到最大尝试次数 {max_attempts}，任务失败。")


def extend_prompt(prompt, max_attempts=5):
    attempts = 0
    while attempts < max_attempts:
        try:
            res = prompt_writer.run(prompt, stream=False)
            logger.info(f"extend prompt finished: {prompt} -> {res.prompt}.")
            return res.prompt
        except Exception as e:
            print(f"尝试 {attempts + 1} 失败：{e}")
            attempts += 1
            if attempts < 5:
                time.sleep(1)  # 延迟重试
    raise gr.Error(f"提示词扩写达到最大尝试次数 {max_attempts}，任务失败。")


if __name__ == '__main__':
    # test_trans()
    pass