moranzhi
bc130d98f4
Introduce WorkflowEngine with state machine, tool registry, and builtin chat template; migrate stream chat to engine callbacks. Move page mode switching to TopBar actions cluster as a ThemeToggle-style dropdown (聊天/工作室/爽文/房间). Co-authored-by: Cursor <cursoragent@cursor.com>
262 lines
9.9 KiB
Python
262 lines
9.9 KiB
Python
"""
|
|
Chat workflow tools extracted from ChatWorkflowService.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
import uuid
|
|
from typing import Any, Dict, List
|
|
|
|
try:
|
|
from backend.models.agent import TurnContext
|
|
from backend.models.internal import CharacterCard, TokenUsageStatus
|
|
from backend.models.regex_rules import RegexPlacement
|
|
from backend.services.character_service import CharacterService
|
|
from backend.services.regex_service import regex_service
|
|
from backend.services.task_queue_manager import TaskType, task_queue_manager
|
|
from backend.services.token_usage_service import token_usage_service
|
|
from backend.core.config import settings
|
|
except ImportError:
|
|
from models.agent import TurnContext
|
|
from models.internal import CharacterCard, TokenUsageStatus
|
|
from models.regex_rules import RegexPlacement
|
|
from services.character_service import CharacterService
|
|
from services.regex_service import regex_service
|
|
from services.task_queue_manager import TaskType, task_queue_manager
|
|
from services.token_usage_service import token_usage_service
|
|
from core.config import settings
|
|
|
|
|
|
_character_service = CharacterService()
|
|
_workflow_service = None
|
|
|
|
|
|
def _get_workflow_service():
|
|
"""Lazy init to avoid circular import with chat_workflow_service."""
|
|
global _workflow_service
|
|
if _workflow_service is None:
|
|
try:
|
|
from backend.services.chat_workflow_service import ChatWorkflowService
|
|
except ImportError:
|
|
from services.chat_workflow_service import ChatWorkflowService
|
|
_workflow_service = ChatWorkflowService()
|
|
return _workflow_service
|
|
|
|
|
|
async def regex_apply_user_input(ctx: TurnContext) -> None:
|
|
processed = regex_service.apply_rules_by_placement(
|
|
text=ctx.user_message,
|
|
placement=RegexPlacement.USER_INPUT.value,
|
|
character_name=ctx.current_role,
|
|
preset_name=ctx.preset_name,
|
|
message_depth=0,
|
|
is_for_llm=True,
|
|
is_markdown_rendered=False,
|
|
)
|
|
if processed != ctx.user_message:
|
|
print("[WorkflowTool] Applied user-input regex rules")
|
|
ctx.user_message = processed
|
|
|
|
|
|
async def load_character(ctx: TurnContext) -> None:
|
|
character_data = ctx.request_data.get("characterData")
|
|
if not character_data:
|
|
character = _character_service.get_character_by_name(ctx.current_role)
|
|
if not character:
|
|
raise ValueError(f"角色 '{ctx.current_role}' 不存在")
|
|
else:
|
|
character = CharacterCard(**character_data)
|
|
ctx.character = character
|
|
print(f"[WorkflowTool] Loaded character: {character.name}")
|
|
|
|
|
|
async def activate_worldbook(ctx: TurnContext) -> None:
|
|
svc = _get_workflow_service()
|
|
active_entries = await svc._collect_and_activate_worldbooks(
|
|
ctx.request_data,
|
|
ctx.character,
|
|
)
|
|
ctx.active_entries = active_entries
|
|
print(f"[WorkflowTool] Activated {len(active_entries)} worldbook entries")
|
|
|
|
if ctx.callbacks and ctx.callbacks.on_worldbook_active:
|
|
entries_payload = [
|
|
entry.model_dump() if hasattr(entry, "model_dump") else entry
|
|
for entry in active_entries
|
|
]
|
|
await ctx.callbacks.on_worldbook_active(entries_payload)
|
|
|
|
|
|
async def load_chat_history(ctx: TurnContext) -> None:
|
|
svc = _get_workflow_service()
|
|
chat_history = await svc._load_chat_history(
|
|
ctx.current_role,
|
|
ctx.current_chat,
|
|
)
|
|
ctx.chat_history = chat_history
|
|
print(f"[WorkflowTool] Loaded {len(chat_history)} history messages")
|
|
|
|
|
|
async def build_prompt_messages(ctx: TurnContext) -> None:
|
|
svc = _get_workflow_service()
|
|
prompt_messages = svc._assemble_prompt(
|
|
ctx.character,
|
|
ctx.chat_history,
|
|
ctx.user_message,
|
|
ctx.active_entries,
|
|
ctx.request_data,
|
|
)
|
|
ctx.prompt_messages = prompt_messages
|
|
print(f"[WorkflowTool] Built {len(prompt_messages)} prompt messages")
|
|
|
|
|
|
async def llm_main_reply(ctx: TurnContext) -> None:
|
|
svc = _get_workflow_service()
|
|
api_config = ctx.request_data.get("apiConfig", {})
|
|
preset_config = ctx.request_data.get("presetConfig", {})
|
|
|
|
if ctx.stream:
|
|
if not api_config.get("api_key"):
|
|
raise ValueError("API Key 未配置，请先在 API 配置页面保存密钥")
|
|
|
|
generated_content = ""
|
|
chunk_count = 0
|
|
start_time = time.time()
|
|
|
|
async for chunk_dict in svc.llm_client.stream_chat(
|
|
messages=ctx.prompt_messages,
|
|
api_url=api_config.get("api_url", ""),
|
|
api_key=api_config.get("api_key", ""),
|
|
model=api_config.get("model", ""),
|
|
temperature=preset_config.get("parameters", {}).get("temperature", 1.0),
|
|
max_tokens=preset_config.get("parameters", {}).get("max_tokens", 30000),
|
|
request_timeout=preset_config.get("parameters", {}).get("request_timeout", 60),
|
|
):
|
|
if isinstance(chunk_dict, dict):
|
|
if chunk_dict.get("type") == "chunk":
|
|
chunk_content = chunk_dict.get("content", "")
|
|
elif chunk_dict.get("type") == "usage":
|
|
continue
|
|
else:
|
|
chunk_content = chunk_dict.get("content", str(chunk_dict))
|
|
else:
|
|
chunk_content = str(chunk_dict)
|
|
|
|
generated_content += chunk_content
|
|
chunk_count += 1
|
|
|
|
if ctx.callbacks and ctx.callbacks.on_chunk:
|
|
await ctx.callbacks.on_chunk(chunk_content)
|
|
|
|
ctx.duration = time.time() - start_time
|
|
ctx.generated_content = generated_content
|
|
ctx.token_usage = {
|
|
"prompt_tokens": len(str(ctx.prompt_messages)) // 4,
|
|
"completion_tokens": len(generated_content) // 4,
|
|
"total_tokens": (len(str(ctx.prompt_messages)) // 4)
|
|
+ (len(generated_content) // 4),
|
|
}
|
|
print(
|
|
f"[WorkflowTool] Stream LLM complete: {chunk_count} chunks, "
|
|
f"{len(generated_content)} chars"
|
|
)
|
|
else:
|
|
result = await svc._generate_response(
|
|
ctx.prompt_messages,
|
|
api_config,
|
|
preset_config,
|
|
stream=False,
|
|
)
|
|
ctx.generated_content = result["content"]
|
|
ctx.token_usage = result.get("usage", {})
|
|
ctx.duration = result.get("duration", 0.0)
|
|
print(f"[WorkflowTool] LLM complete: {len(ctx.generated_content)} chars")
|
|
|
|
|
|
async def regex_apply_ai_output(ctx: TurnContext) -> None:
|
|
processed = regex_service.apply_rules_by_placement(
|
|
text=ctx.generated_content,
|
|
placement=RegexPlacement.AI_OUTPUT.value,
|
|
character_name=ctx.current_role,
|
|
preset_name=ctx.preset_name,
|
|
message_depth=0,
|
|
is_for_llm=False,
|
|
is_markdown_rendered=False,
|
|
)
|
|
if processed != ctx.generated_content:
|
|
print("[WorkflowTool] Applied AI-output regex rules")
|
|
ctx.generated_content = processed
|
|
|
|
|
|
async def record_token_usage(ctx: TurnContext) -> None:
|
|
chat_id = f"{ctx.current_role}/{ctx.current_chat}"
|
|
floor = ctx.request_data.get("floor", 0)
|
|
api_config = ctx.request_data.get("apiConfig", {})
|
|
|
|
try:
|
|
await token_usage_service.record_usage(
|
|
chat_id=chat_id,
|
|
role_name=ctx.current_role,
|
|
chat_name=ctx.current_chat,
|
|
prompt_tokens=ctx.token_usage.get("prompt_tokens", 0),
|
|
completion_tokens=ctx.token_usage.get("completion_tokens", 0),
|
|
total_tokens=ctx.token_usage.get("total_tokens", 0),
|
|
status=TokenUsageStatus.COMPLETED,
|
|
floor=floor + 1,
|
|
duration=ctx.duration,
|
|
model=api_config.get("model"),
|
|
api_provider="openai",
|
|
api_url=api_config.get("api_url"),
|
|
)
|
|
except Exception as exc:
|
|
print(f"[WorkflowTool] Token usage recording failed: {exc}")
|
|
|
|
|
|
async def enqueue_parallel_tasks(ctx: TurnContext) -> None:
|
|
chat_id = f"{ctx.current_role}/{ctx.current_chat}"
|
|
options = ctx.request_data.get("options", {})
|
|
image_task_id = None
|
|
table_task_id = None
|
|
|
|
if options.get("imageWorkflow", False):
|
|
image_task_id = f"img_{uuid.uuid4().hex[:8]}"
|
|
await task_queue_manager.add_task(image_task_id, TaskType.IMAGE_WORKFLOW, chat_id)
|
|
|
|
if options.get("dynamicTable", False):
|
|
table_task_id = f"tbl_{uuid.uuid4().hex[:8]}"
|
|
await task_queue_manager.add_task(table_task_id, TaskType.DYNAMIC_TABLE, chat_id)
|
|
|
|
ctx.task_ids = {
|
|
"imageWorkflow": image_task_id,
|
|
"dynamicTable": table_task_id,
|
|
}
|
|
|
|
if ctx.callbacks and ctx.callbacks.on_tasks_created:
|
|
if image_task_id or table_task_id:
|
|
await ctx.callbacks.on_tasks_created(ctx.task_ids)
|
|
|
|
# Fire-and-forget parallel workers (same as legacy service)
|
|
svc = _get_workflow_service()
|
|
asyncio.create_task(
|
|
svc._start_parallel_tasks(
|
|
ctx.request_data,
|
|
ctx.generated_content,
|
|
image_task_id,
|
|
table_task_id,
|
|
)
|
|
)
|
|
|
|
|
|
def register_chat_tools(registry) -> None:
|
|
"""Register all chat workflow tools on the given registry."""
|
|
registry.register("regex_apply_user_input", regex_apply_user_input, description="Apply user-input regex")
|
|
registry.register("load_character", load_character, description="Load character card")
|
|
registry.register("activate_worldbook", activate_worldbook, description="Activate worldbook entries")
|
|
registry.register("load_chat_history", load_chat_history, description="Load chat history")
|
|
registry.register("build_prompt_messages", build_prompt_messages, description="Assemble LLM prompt")
|
|
registry.register("llm_main_reply", llm_main_reply, description="Call main LLM (supports stream)")
|
|
registry.register("regex_apply_ai_output", regex_apply_ai_output, description="Apply AI-output regex")
|
|
registry.register("record_token_usage", record_token_usage, description="Persist token usage")
|
|
registry.register("enqueue_parallel_tasks", enqueue_parallel_tasks, description="Enqueue parallel tasks")
|