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