open-webui-addons

Story-Router-helper.py

"""
title: Story Route Helper
author: deepseek,luojia
author_url: https://luojia.me
version: 0.5.0
description: 此脚本用于剧情路线交互时处理消息，提供两种处理模式：合并剧情或仅清除选项。用法及System prompt建议见 https://luojia.me/?p=10016
"""

from pydantic import BaseModel, Field
from typing import Optional
import re


class Filter:
    class Valves(BaseModel):
        priority: int = Field(default=10, description="过滤器操作的优先级级别。")
        max_turns: int = Field(default=20, description="用户允许的最大对话轮数。")
        merge_assistant_messages: bool = Field(
            default=True,
            description="True=合并所有助理消息并精简结构；False=仅移除选项保留原有消息结构",
        )

    class UserValves(BaseModel):
        max_turns: int = Field(default=20, description="用户允许的最大对话轮数。")

    def __init__(self):
        self.valves = self.Valves()

    def _is_story_route_mode(self, messages) -> bool:
        return (
            len(messages) > 1
            and messages[1].get("role") == "user"
            and messages[1].get("content", "").startswith("分支剧情场景定义：")
        )

    def inlet(self, body: dict, __user__: Optional[dict] = None) -> dict:
        messages = body.get("messages", [])
        if not messages or messages[0].get("role") != "system":
            return body
        if not isinstance(messages[1].get("content"), str):
            return body

        # 调试输出原始消息
        print("\n=============== 原始消息 ===============")
        for i, msg in enumerate(messages):
            print(f"[{i}][{msg['role']}]: {msg.get('content','')[:80]}...")

        if not self._is_story_route_mode(messages):
            return body

        print(
            "\n分支剧情模式激活，处理方式:",
            "合并消息" if self.valves.merge_assistant_messages else "仅清除选项",
        )

        if self.valves.merge_assistant_messages:
            # 模式1：合并消息的精简结构
            system_msg = messages[0]
            first_user_msg = messages[1]
            last_user_msg = None
            merged_content = []

            for msg in messages[2:]:
                if msg["role"] == "user":
                    last_user_msg = msg
                elif msg["role"] == "assistant":
                    cleaned = re.sub(
                        r"```options.*?```", "", msg.get("content", ""), flags=re.DOTALL
                    ).strip()
                    if cleaned:
                        merged_content.append(cleaned)

            new_messages = [system_msg, first_user_msg]
            if merged_content:
                new_messages.append(
                    {"role": "assistant", "content": "\n\n".join(merged_content)}
                )
            if last_user_msg:
                new_messages.append(last_user_msg)

            body["messages"] = new_messages
        else:
            # 模式2：保留结构的选项清除
            for msg in messages:
                if msg["role"] == "assistant":
                    msg["content"] = re.sub(
                        r"```options.*?```", "", msg.get("content", ""), flags=re.DOTALL
                    ).strip()
            body["messages"] = messages

        # 调试输出处理结果
        print("\n=============== 处理后的消息 ===============")
        for i, msg in enumerate(body["messages"]):
            print(f"[{i}][{msg['role']}]: {msg.get('content','')[:80]}...")
        print("=========================================\n")

        return body

    def outlet(self, body: dict, __user__: Optional[dict] = None) -> dict:
        return body