145 lines
6.5 KiB
Python
145 lines
6.5 KiB
Python
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---"
|
|
|