from __future__ import annotations import json import os from dataclasses import dataclass from functools import lru_cache from typing import Any, Dict, List, Sequence, Set from .manifest_loader import ManifestPromptLoader @dataclass class PromptPackage: system_prompt: str user_augmentation: str allowed_nodes: Dict[str, List[str]] class PromptComposer: def __init__(self, prompts_dir: str): self.loader = ManifestPromptLoader(prompts_dir) self.prompts_dir = prompts_dir self._simple_prompt = self.loader.load_text_file("simple_mode_prompt.txt") @lru_cache(maxsize=64) def _cached_macro_scene_parts(self, drone_state: str) -> tuple[str, str]: header = self._load_partial("macro_header.txt") template = self._load_partial("template_ground.txt" if drone_state == "on_ground" else "template_air.txt") return header, template @lru_cache(maxsize=64) def _cached_micro_scene_parts(self) -> str: return self._load_partial("micro_header.txt") def _load_partial(self, file_name: str) -> str: path = os.path.join(self.prompts_dir, "partials", file_name) if not os.path.exists(path): return "" return self.loader.load_fragment_text(path).strip("\n") def _load_nodes_payload(self) -> Dict[str, List[Dict[str, str]]]: core_path = os.path.join(self.prompts_dir, "partials", "core_nodes.json") if not os.path.exists(core_path): return {"actions": [], "conditions": [], "control_flow": [], "decorators": []} with open(core_path, "r", encoding="utf-8") as f: payload = json.load(f) return payload def _slice_nodes(self, required_actions: Sequence[str], risk_flags: Sequence[str], scene_mode: str) -> Dict[str, List[str]]: payload = self._load_nodes_payload() base_actions: Set[str] = {"takeoff", "land", "fly_to_waypoint", "move_direction"} base_conditions: Set[str] = {"object_detected"} if scene_mode != "simple": base_actions.update({"rotate_search", "take_photos"}) action_names = {a.get("name") for a in payload.get("actions", []) if isinstance(a, dict)} condition_names = {c.get("name") for c in payload.get("conditions", []) if isinstance(c, dict)} selected_actions = set(required_actions) | (base_actions & action_names) selected_conditions = base_conditions & condition_names if "needs_manual_confirmation" in risk_flags and "manual_confirmation" in action_names: selected_actions.add("manual_confirmation") if "needs_emergency_return_rule" in risk_flags and "return_emergency" in action_names: selected_actions.add("return_emergency") if "needs_approach_target" in risk_flags and "approach_target" in action_names: selected_actions.add("approach_target") ordered_actions = sorted(selected_actions)[:30] ordered_conditions = sorted(selected_conditions)[:30] return {"actions": ordered_actions, "conditions": ordered_conditions} def _build_nodes_snippet(self, selected: Dict[str, List[str]]) -> str: payload = self._load_nodes_payload() selected_action_set = set(selected.get("actions", [])) selected_condition_set = set(selected.get("conditions", [])) slim_payload = { "actions": [a for a in payload.get("actions", []) if a.get("name") in selected_action_set], "conditions": [c for c in payload.get("conditions", []) if c.get("name") in selected_condition_set], "control_flow": payload.get("control_flow", []), "decorators": payload.get("decorators", []), } return "\n".join( [ "## 一、核心节点定义（裁剪后）", "#### 1. 可用节点定义 (必须遵守)", "```json", json.dumps(slim_payload, ensure_ascii=False, indent=2), "```", ] ) def compose_macro( self, scene_mode: str, drone_state: str, intent_type: str, required_actions: Sequence[str], risk_flags: Sequence[str], context_blocks: Dict[str, str], include_extra_examples: bool, ) -> PromptPackage: if scene_mode == "simple": user_aug = self._build_user_augmentation(context_blocks) return PromptPackage(system_prompt=self._simple_prompt, user_augmentation=user_aug, allowed_nodes={}) header, template = self._cached_macro_scene_parts(drone_state) selected_nodes = self._slice_nodes(required_actions, risk_flags, scene_mode) node_snippet = self._build_nodes_snippet(selected_nodes) parts: List[str] = [header, node_snippet, template] common_rules = self._load_partial("common_rules.txt") if common_rules: parts.append(common_rules) if include_extra_examples: extra = self._load_partial("system_extra_examples.txt") if extra: parts.append(extra) parts.append(f"## 任务意图标签\n- intent_type: `{intent_type}`") system_prompt = "\n\n".join(p for p in parts if p).strip() return PromptPackage(system_prompt=system_prompt, user_augmentation=self._build_user_augmentation(context_blocks), allowed_nodes=selected_nodes) def compose_micro(self, macro_tree: Dict[str, Any], resolved_data: Dict[str, Any], atomic_schema: Dict[str, Any]) -> str: header = self._cached_micro_scene_parts() parts = [ header, "## 1. 原宏观骨架树 (macro_tree)", "```json\n" + json.dumps(macro_tree, ensure_ascii=False, indent=2) + "\n```", "## 2. 确切数据字典 (resolved_data)", "```json\n" + json.dumps(resolved_data, ensure_ascii=False, indent=2) + "\n```", "## 3. 原子节点规范 (atomic_schema)", "```json\n" + json.dumps(atomic_schema, ensure_ascii=False, indent=2) + "\n```", ] return "\n\n".join(parts) def _build_user_augmentation(self, context_blocks: Dict[str, str]) -> str: ordered = [("地点知识", "location"), ("任务模式", "pattern"), ("规则知识", "rules")] chunks: List[str] = [] for label, key in ordered: value = (context_blocks.get(key) or "").strip() if value: chunks.append(f"【{label}】\n{value}") if not chunks: return "" return "\n\n---\n参考知识：\n" + "\n\n".join(chunks) + "\n---"