|
@@ -0,0 +1,305 @@
|
|
|
|
|
+"""AI Plan Generator (P3-M3).
|
|
|
|
|
+
|
|
|
|
|
+Converts natural language requirements into structured simulation plans
|
|
|
|
|
+using Kimi k3 model, with L0 pre-screening validation.
|
|
|
|
|
+"""
|
|
|
|
|
+import json
|
|
|
|
|
+from pathlib import Path
|
|
|
|
|
+from typing import Dict, List, Optional, Any, Tuple
|
|
|
|
|
+
|
|
|
|
|
+from ..config import PROMPTS_DIR
|
|
|
|
|
+from ..services.ai_client import get_kimi_client
|
|
|
|
|
+from ..services.l0_prescreening import L0PreScreeningEngine
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+class AIPlanGenerator:
|
|
|
|
|
+ """Generate simulation plans from natural language using AI."""
|
|
|
|
|
+
|
|
|
|
|
+ def __init__(self, l0_engine: Optional[L0PreScreeningEngine] = None):
|
|
|
|
|
+ self.ai_client = get_kimi_client()
|
|
|
|
|
+ self.l0_engine = l0_engine or L0PreScreeningEngine()
|
|
|
|
|
+ self._prompt_template = None
|
|
|
|
|
+
|
|
|
|
|
+ def _load_prompt(self) -> str:
|
|
|
|
|
+ """Load plan generation prompt template."""
|
|
|
|
|
+ if self._prompt_template is None:
|
|
|
|
|
+ prompt_path = PROMPTS_DIR / "plan_generation" / "generate.txt"
|
|
|
|
|
+ if prompt_path.exists():
|
|
|
|
|
+ with open(prompt_path, "r", encoding="utf-8") as f:
|
|
|
|
|
+ self._prompt_template = f.read()
|
|
|
|
|
+ else:
|
|
|
|
|
+ self._prompt_template = self._default_prompt()
|
|
|
|
|
+ return self._prompt_template
|
|
|
|
|
+
|
|
|
|
|
+ def _default_prompt(self) -> str:
|
|
|
|
|
+ """Fallback default prompt."""
|
|
|
|
|
+ return """你是轴向磁通电机仿真方案专家。将用户需求转化为JSON格式的仿真方案,包含plan_name、topology、boundary_conditions、scan_variables、search_strategy、acceptance_criteria、reasoning。输出纯JSON。"""
|
|
|
|
|
+
|
|
|
|
|
+ def generate(
|
|
|
|
|
+ self,
|
|
|
|
|
+ user_requirement: str,
|
|
|
|
|
+ project_context: Optional[Dict[str, Any]] = None,
|
|
|
|
|
+ existing_experience: Optional[List[Dict[str, Any]]] = None,
|
|
|
|
|
+ ) -> Dict[str, Any]:
|
|
|
|
|
+ """Generate a simulation plan from natural language.
|
|
|
|
|
+
|
|
|
|
|
+ Args:
|
|
|
|
|
+ user_requirement: Natural language description of the simulation need.
|
|
|
|
|
+ project_context: Optional project context (topology, existing boundary conditions).
|
|
|
|
|
+ existing_experience: Optional similar experience cases for reference.
|
|
|
|
|
+
|
|
|
|
|
+ Returns:
|
|
|
|
|
+ Dict with generated plan, validation results, and AI reasoning.
|
|
|
|
|
+ """
|
|
|
|
|
+ if not self.ai_client.is_configured:
|
|
|
|
|
+ raise RuntimeError("KIMI_API_KEY is not configured")
|
|
|
|
|
+
|
|
|
|
|
+ # Build context message
|
|
|
|
|
+ context_parts = []
|
|
|
|
|
+ if project_context:
|
|
|
|
|
+ context_parts.append(f"项目上下文:{json.dumps(project_context, ensure_ascii=False)}")
|
|
|
|
|
+ if existing_experience:
|
|
|
|
|
+ context_parts.append(f"参考经验案例({len(existing_experience)}个):{json.dumps(existing_experience[:3], ensure_ascii=False)}")
|
|
|
|
|
+
|
|
|
|
|
+ user_message = user_requirement
|
|
|
|
|
+ if context_parts:
|
|
|
|
|
+ user_message += "\n\n" + "\n".join(context_parts)
|
|
|
|
|
+
|
|
|
|
|
+ # Call AI
|
|
|
|
|
+ system_prompt = self._load_prompt()
|
|
|
|
|
+ result = self.ai_client.chat_json(
|
|
|
|
|
+ messages=[{"role": "user", "content": user_message}],
|
|
|
|
|
+ system_prompt=system_prompt,
|
|
|
|
|
+ max_tokens=4000,
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ # Parse generated plan
|
|
|
|
|
+ plan = result.get("parsed_json", {})
|
|
|
|
|
+ if not plan:
|
|
|
|
|
+ # Try to extract and repair JSON from raw content
|
|
|
|
|
+ raw = result.get("raw_content", result.get("content", ""))
|
|
|
|
|
+ plan = self._try_repair_json(raw)
|
|
|
|
|
+ if not plan:
|
|
|
|
|
+ plan = {"error": "Failed to parse AI response", "raw_content": raw[:1000]}
|
|
|
|
|
+
|
|
|
|
|
+ # Validate with L0 pre-screening
|
|
|
|
|
+ validation = self._validate_plan(plan)
|
|
|
|
|
+
|
|
|
|
|
+ return {
|
|
|
|
|
+ "plan": plan,
|
|
|
|
|
+ "validation": validation,
|
|
|
|
|
+ "ai_reasoning": plan.get("reasoning", ""),
|
|
|
|
|
+ "usage": result.get("usage", {}),
|
|
|
|
|
+ "raw_content": result.get("content", ""),
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ def _try_repair_json(self, raw: str) -> Dict[str, Any]:
|
|
|
|
|
+ """Try to extract and repair JSON from potentially truncated AI output.
|
|
|
|
|
+
|
|
|
|
|
+ Handles:
|
|
|
|
|
+ - Markdown code fences
|
|
|
|
|
+ - Truncated JSON (missing closing braces)
|
|
|
|
|
+ - Extra text before/after JSON
|
|
|
|
|
+ """
|
|
|
|
|
+ if not raw:
|
|
|
|
|
+ return {}
|
|
|
|
|
+
|
|
|
|
|
+ text = raw.strip()
|
|
|
|
|
+
|
|
|
|
|
+ # Remove markdown code fences
|
|
|
|
|
+ if text.startswith("```"):
|
|
|
|
|
+ lines = text.split("\n")
|
|
|
|
|
+ if lines[0].startswith("```"):
|
|
|
|
|
+ lines = lines[1:]
|
|
|
|
|
+ if lines and lines[-1].strip() == "```":
|
|
|
|
|
+ lines = lines[:-1]
|
|
|
|
|
+ text = "\n".join(lines).strip()
|
|
|
|
|
+
|
|
|
|
|
+ # Find first { and last }
|
|
|
|
|
+ first_brace = text.find("{")
|
|
|
|
|
+ last_brace = text.rfind("}")
|
|
|
|
|
+
|
|
|
|
|
+ if first_brace == -1:
|
|
|
|
|
+ return {}
|
|
|
|
|
+
|
|
|
|
|
+ if last_brace == -1 or last_brace < first_brace:
|
|
|
|
|
+ # JSON is truncated - try to close it
|
|
|
|
|
+ json_str = text[first_brace:]
|
|
|
|
|
+ # Count open braces and close them
|
|
|
|
|
+ open_braces = json_str.count("{") - json_str.count("}")
|
|
|
|
|
+ open_brackets = json_str.count("[") - json_str.count("]")
|
|
|
|
|
+ # Remove trailing incomplete content
|
|
|
|
|
+ last_comma = max(json_str.rfind(","), json_str.rfind(":"))
|
|
|
|
|
+ if last_comma > len(json_str) * 0.8:
|
|
|
|
|
+ json_str = json_str[:last_comma]
|
|
|
|
|
+ # Close brackets and braces
|
|
|
|
|
+ json_str += "]" * max(0, open_brackets)
|
|
|
|
|
+ json_str += "}" * max(0, open_braces)
|
|
|
|
|
+ else:
|
|
|
|
|
+ json_str = text[first_brace:last_brace + 1]
|
|
|
|
|
+
|
|
|
|
|
+ try:
|
|
|
|
|
+ return json.loads(json_str)
|
|
|
|
|
+ except json.JSONDecodeError:
|
|
|
|
|
+ # Try one more repair: remove trailing comma before closing
|
|
|
|
|
+ try:
|
|
|
|
|
+ repaired = json_str.rstrip()
|
|
|
|
|
+ while repaired and repaired[-1] in ", \n\t":
|
|
|
|
|
+ repaired = repaired[:-1]
|
|
|
|
|
+ # Re-count and close
|
|
|
|
|
+ open_braces = repaired.count("{") - repaired.count("}")
|
|
|
|
|
+ open_brackets = repaired.count("[") - repaired.count("]")
|
|
|
|
|
+ repaired += "]" * max(0, open_brackets)
|
|
|
|
|
+ repaired += "}" * max(0, open_braces)
|
|
|
|
|
+ return json.loads(repaired)
|
|
|
|
|
+ except (json.JSONDecodeError, Exception):
|
|
|
|
|
+ return {}
|
|
|
|
|
+
|
|
|
|
|
+ def _validate_plan(self, plan: Dict[str, Any]) -> Dict[str, Any]:
|
|
|
|
|
+ """Validate generated plan with L0 pre-screening.
|
|
|
|
|
+
|
|
|
|
|
+ Checks:
|
|
|
|
|
+ 1. Boundary conditions feasibility
|
|
|
|
|
+ 2. Scan variable ranges within engineering limits
|
|
|
|
|
+ 3. Sample point count estimation
|
|
|
|
|
+ 4. Required fields presence
|
|
|
|
|
+ """
|
|
|
|
|
+ issues = []
|
|
|
|
|
+ warnings = []
|
|
|
|
|
+ checks = []
|
|
|
|
|
+
|
|
|
|
|
+ # Check required fields
|
|
|
|
|
+ required_fields = ["plan_name", "topology", "scan_variables"]
|
|
|
|
|
+ for field in required_fields:
|
|
|
|
|
+ if field not in plan:
|
|
|
|
|
+ issues.append(f"Missing required field: {field}")
|
|
|
|
|
+
|
|
|
|
|
+ # Validate topology
|
|
|
|
|
+ topology = plan.get("topology", "")
|
|
|
|
|
+ if topology and topology not in ("SSSR", "DRSS", "SDSR"):
|
|
|
|
|
+ warnings.append(f"Unusual topology: {topology}")
|
|
|
|
|
+
|
|
|
|
|
+ # Validate boundary conditions with L0
|
|
|
|
|
+ bc = plan.get("boundary_conditions", {})
|
|
|
|
|
+ if bc:
|
|
|
|
|
+ l0_report = self.l0_engine.evaluate(bc)
|
|
|
|
|
+ checks.append({
|
|
|
|
|
+ "name": "boundary_conditions_l0",
|
|
|
|
|
+ "feasible": l0_report.feasible,
|
|
|
|
|
+ "passed": l0_report.passed_checks,
|
|
|
|
|
+ "total": l0_report.total_checks,
|
|
|
|
|
+ "failed_items": [r.name for r in l0_report.results if not r.passed],
|
|
|
|
|
+ })
|
|
|
|
|
+ if not l0_report.feasible:
|
|
|
|
|
+ issues.append(f"Boundary conditions infeasible: {l0_report.failed_checks} checks failed")
|
|
|
|
|
+
|
|
|
|
|
+ # Validate scan variables
|
|
|
|
|
+ variables = plan.get("scan_variables", [])
|
|
|
|
|
+ if not isinstance(variables, list):
|
|
|
|
|
+ variables = []
|
|
|
|
|
+ total_points = 1
|
|
|
|
|
+ normalized_vars = []
|
|
|
|
|
+ for var in variables:
|
|
|
|
|
+ if isinstance(var, str):
|
|
|
|
|
+ var = {"name": var, "min_value": None, "max_value": None}
|
|
|
|
|
+ if not isinstance(var, dict):
|
|
|
|
|
+ continue
|
|
|
|
|
+ normalized_vars.append(var)
|
|
|
|
|
+ name = var.get("name", "unknown")
|
|
|
|
|
+ min_v = var.get("min_value")
|
|
|
|
|
+ max_v = var.get("max_value")
|
|
|
|
|
+ step = var.get("step")
|
|
|
|
|
+
|
|
|
|
|
+ if min_v is None or max_v is None:
|
|
|
|
|
+ issues.append(f"Variable {name}: missing min/max values")
|
|
|
|
|
+ continue
|
|
|
|
|
+
|
|
|
|
|
+ if min_v >= max_v:
|
|
|
|
|
+ issues.append(f"Variable {name}: min ({min_v}) >= max ({max_v})")
|
|
|
|
|
+
|
|
|
|
|
+ if step and step > 0:
|
|
|
|
|
+ n_points = int((max_v - min_v) / step) + 1
|
|
|
|
|
+ total_points *= n_points
|
|
|
|
|
+ if n_points > 50:
|
|
|
|
|
+ warnings.append(f"Variable {name}: {n_points} levels may be too many")
|
|
|
|
|
+ plan["scan_variables"] = normalized_vars
|
|
|
|
|
+
|
|
|
|
|
+ checks.append({
|
|
|
|
|
+ "name": "scan_variables",
|
|
|
|
|
+ "count": len(variables),
|
|
|
|
|
+ "estimated_full_factorial_points": total_points,
|
|
|
|
|
+ "adaptive_search_recommended": total_points > 100,
|
|
|
|
|
+ })
|
|
|
|
|
+
|
|
|
|
|
+ if total_points > 500:
|
|
|
|
|
+ warnings.append(f"Full factorial would require {total_points} points - strongly recommend adaptive search")
|
|
|
|
|
+
|
|
|
|
|
+ # Validate search strategy
|
|
|
|
|
+ search = plan.get("search_strategy", {})
|
|
|
|
|
+ if search:
|
|
|
|
|
+ max_calls = search.get("max_solver_calls", 80)
|
|
|
|
|
+ if max_calls < total_points and not search.get("method", "").startswith(("constrained", "active")):
|
|
|
|
|
+ warnings.append(f"Budget ({max_calls}) < full factorial ({total_points}) but method is not adaptive")
|
|
|
|
|
+
|
|
|
|
|
+ overall_valid = len(issues) == 0
|
|
|
|
|
+ return {
|
|
|
|
|
+ "valid": overall_valid,
|
|
|
|
|
+ "issues": issues,
|
|
|
|
|
+ "warnings": warnings,
|
|
|
|
|
+ "checks": checks,
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ def refine_plan(
|
|
|
|
|
+ self,
|
|
|
|
|
+ original_plan: Dict[str, Any],
|
|
|
|
|
+ user_feedback: str,
|
|
|
|
|
+ ) -> Dict[str, Any]:
|
|
|
|
|
+ """Refine an existing plan based on user feedback.
|
|
|
|
|
+
|
|
|
|
|
+ Args:
|
|
|
|
|
+ original_plan: The previously generated plan.
|
|
|
|
|
+ user_feedback: Natural language feedback for refinement.
|
|
|
|
|
+
|
|
|
|
|
+ Returns:
|
|
|
|
|
+ Refined plan with validation.
|
|
|
|
|
+ """
|
|
|
|
|
+ if not self.ai_client.is_configured:
|
|
|
|
|
+ raise RuntimeError("KIMI_API_KEY is not configured")
|
|
|
|
|
+
|
|
|
|
|
+ system_prompt = self._load_prompt() + "\n\n你正在优化一个已有的仿真方案。根据用户反馈调整方案,保持其他部分不变。"
|
|
|
|
|
+
|
|
|
|
|
+ user_message = f"原有方案:\n{json.dumps(original_plan, ensure_ascii=False, indent=2)}\n\n用户反馈:{user_feedback}\n\n请输出优化后的完整方案JSON。"
|
|
|
|
|
+
|
|
|
|
|
+ result = self.ai_client.chat_json(
|
|
|
|
|
+ messages=[{"role": "user", "content": user_message}],
|
|
|
|
|
+ system_prompt=system_prompt,
|
|
|
|
|
+ max_tokens=4000,
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ refined_plan = result.get("parsed_json", {})
|
|
|
|
|
+ if not refined_plan:
|
|
|
|
|
+ raw = result.get("raw_content", result.get("content", ""))
|
|
|
|
|
+ refined_plan = self._try_repair_json(raw)
|
|
|
|
|
+ if not refined_plan:
|
|
|
|
|
+ refined_plan = {"error": "Failed to parse AI response", "raw_content": raw[:1000]}
|
|
|
|
|
+
|
|
|
|
|
+ validation = self._validate_plan(refined_plan)
|
|
|
|
|
+
|
|
|
|
|
+ return {
|
|
|
|
|
+ "plan": refined_plan,
|
|
|
|
|
+ "validation": validation,
|
|
|
|
|
+ "ai_reasoning": refined_plan.get("reasoning", ""),
|
|
|
|
|
+ "usage": result.get("usage", {}),
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+# Global singleton
|
|
|
|
|
+_generator: Optional[AIPlanGenerator] = None
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def get_plan_generator() -> AIPlanGenerator:
|
|
|
|
|
+ """Get or create global AIPlanGenerator singleton."""
|
|
|
|
|
+ global _generator
|
|
|
|
|
+ if _generator is None:
|
|
|
|
|
+ _generator = AIPlanGenerator()
|
|
|
|
|
+ return _generator
|