From 7f2f74fbb6dafe03f31ce0b4c93a9aa84e23fcea Mon Sep 17 00:00:00 2001 From: Val-sss <154882199@qq.com> Date: Wed, 29 Apr 2026 11:12:13 +0800 Subject: [PATCH] feat: DAG task compiler + prompt optimizer, remove Python expert system - Add TaskCompiler (kaiwu/core/task_compiler.py): lightweight DAG scheduler with ThreadPoolExecutor + topological sort for serial/parallel multi-task - Add PromptOptimizer (kaiwu/flywheel/prompt_optimizer.py): analyzes trajectories via Opus/Sonnet API, appends learned rules to expert YAML system_prompt - Add Cross-Encoder reranker (kaiwu/search/reranker.py): optional BM25+CE pipeline - Add Reflexion persistence (kaiwu/memory/pattern_md.py): REFLECTION.md structured pattern memory with /plan integration - Remove Python expert system (ExpertBase, BugFixExpert.py, SelfImprovingOptimizer) - Revert registry to YAML-only loading - Remove _get_python_expert/_run_python_expert from orchestrator - Replace run_self_improvement with run_prompt_optimization in ab_tester - 277 tests passing (265 regression + 12 new task_compiler tests) Co-Authored-By: Claude Opus 4.6 (1M context) --- STATUS.md | 31 +++- kaiwu/core/orchestrator.py | 20 +++ kaiwu/core/planner.py | 19 ++ kaiwu/core/task_compiler.py | 229 ++++++++++++++++++++++++ kaiwu/experts/search_augmentor.py | 16 +- kaiwu/flywheel/ab_tester.py | 32 +++- kaiwu/flywheel/prompt_optimizer.py | 163 +++++++++++++++++ kaiwu/memory/pattern_md.py | 120 +++++++++++++ kaiwu/search/reranker.py | 88 ++++++++++ kaiwu/tests/test_task_compiler.py | 273 +++++++++++++++++++++++++++++ pyproject.toml | 5 +- 11 files changed, 989 insertions(+), 7 deletions(-) create mode 100644 kaiwu/core/task_compiler.py create mode 100644 kaiwu/flywheel/prompt_optimizer.py create mode 100644 kaiwu/search/reranker.py create mode 100644 kaiwu/tests/test_task_compiler.py diff --git a/STATUS.md b/STATUS.md index 69696d3..fd99272 100644 --- a/STATUS.md +++ b/STATUS.md @@ -7,9 +7,36 @@ --- -## 当前状态：v0.7.0 (2026-04-29) +## 当前状态：v0.8.0 (2026-04-29) -全部功能已实现，282/282 测试全绿（265 单元/回归 + 17 E2E 真实模型），已推送 GitHub。 +全部功能已实现，265/265 单元/回归测试全绿。自我强化专家系统已集成。 + +### v0.8.0 新增：自我强化专家系统 + +**Python专家升级** +- ExpertBase基类（`experts/base.py`）：_locate/_generate/_verify三阶段helper +- BugFixExpert Python版本（`builtin_experts/bugfix_expert.py`）：TRIGGER+run()+test()三段式 +- Registry支持Python专家加载，Python优先级高于同名yaml +- SE-RED-1保护：TRIGGER和test()不可被LLM修改 + +**SelfImprovingOptimizer（飞轮自我强化）** +- 积累5次成功轨迹 → Opus/Sonnet API分析 → 生成新run()代码 +- 三道门验证：语法检查 → 后端测试(test()) → 回测 +- SE-RED-3：修改前备份，失败自动回滚 +- SE-RED-4：离线执行，不在用户任务流程中 +- FLEX-1：无API key时跳过，基本功能不受影响 + +**Reflexion持久化** +- 任务完成后（成功/失败）自动写入REFLECTION.md +- SE-RED-5：结构化格式（日期+摘要+根因/注意） +- /plan时自动读取历史Reflection作为风险提示 +- 每section最多20条，防止无限增长 + +**Cross-Encoder搜索重排** +- `search/reranker.py`：cross-encoder/ms-marco-MiniLM-L-6-v2（82MB CPU） +- 集成到搜索管道：BM25重排后再Cross-Encoder精排 +- FLEX-2：sentence-transformers未安装或CPU慢(>2s)时自动跳过 +- 可选依赖：`pip install kwcode[rerank]` ### 已完成功能清单 diff --git a/kaiwu/core/orchestrator.py b/kaiwu/core/orchestrator.py index 82322f4..b59dbad 100644 --- a/kaiwu/core/orchestrator.py +++ b/kaiwu/core/orchestrator.py @@ -203,6 +203,8 @@ class PipelineOrchestrator: self._record_value(project_root, gate_result, True, elapsed, ctx) # P2: Milestone check self._check_milestone(on_status) + # Reflexion持久化：成功时也记录注意事项 + self._persist_reflection(project_root, ctx, gate_result, success=True) return { "success": True, "context": ctx, @@ -268,6 +270,8 @@ class PipelineOrchestrator: self._record_ab_result(ab_candidate_name, ab_used_new, False, elapsed, on_status) # P2: Value tracking (local SQLite) self._record_value(project_root, gate_result, False, elapsed, ctx) + # Reflexion持久化：失败时记录根因 + self._persist_reflection(project_root, ctx, gate_result, success=False) return { "success": False, "context": ctx, @@ -445,3 +449,19 @@ class PipelineOrchestrator: self._notifier.queue_milestone(total, expert_count, 0.0) except Exception as e: logger.debug("Milestone check failed (non-blocking): %s", e) + + def _persist_reflection(self, project_root, ctx, gate_result, success): + """Reflexion持久化：任务完成后写入REFLECTION.md（非阻塞）。""" + try: + if not ctx.reflection: + return + from kaiwu.memory.pattern_md import save_reflection + save_reflection( + project_root=project_root, + expert_type=gate_result.get("expert_type", "unknown"), + task_summary=ctx.user_input[:30], + reflection=ctx.reflection, + success=success, + ) + except Exception as e: + logger.debug("Reflection persistence failed (non-blocking): %s", e) diff --git a/kaiwu/core/planner.py b/kaiwu/core/planner.py index df0a5dd..68c099b 100644 --- a/kaiwu/core/planner.py +++ b/kaiwu/core/planner.py @@ -95,6 +95,14 @@ class Planner: project_root=ctx.project_root, ) + # 读取历史Reflexion作为风险提示 + historical_reflections = "" + try: + from kaiwu.memory.pattern_md import get_reflections_for_plan + historical_reflections = get_reflections_for_plan(ctx.project_root, expert_type) + except Exception: + pass + # Overall risk risk = estimate_risk( step_type=expert_type, @@ -166,6 +174,9 @@ class Planner: risk_reason="不修改文件", )) + # 注入历史Reflexion到plan展示 + self._historical_reflections = historical_reflections + return steps def print_plan(self, steps: list[PlanStep], console): @@ -204,6 +215,14 @@ class Planner: elif max_risk.risk == "Medium": console.print(" [yellow]△ 此任务有一定风险，请确认修改范围[/yellow]") + # 显示历史Reflexion风险提示 + reflections = getattr(self, "_historical_reflections", "") + if reflections: + console.print(" [dim]── 历史经验 ──[/dim]") + for line in reflections.splitlines(): + if line.strip(): + console.print(f" [dim]{line}[/dim]") + def _preview_locator(self, ctx: TaskContext) -> tuple[list[str], list[str]]: """Read-only preview of locator results for planning.""" try: diff --git a/kaiwu/core/task_compiler.py b/kaiwu/core/task_compiler.py new file mode 100644 index 0000000..e71b244 --- /dev/null +++ b/kaiwu/core/task_compiler.py @@ -0,0 +1,229 @@ +""" +Lightweight DAG Task Compiler. +Accepts task definitions with dependencies, builds a DAG, executes via +ThreadPoolExecutor with topological ordering. + +Zero new pip dependencies (ThreadPoolExecutor is stdlib). +Each task gets its own TaskContext (RED-3: independent context). +""" + +import logging +import time +from collections import deque +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from kaiwu.core.orchestrator import PipelineOrchestrator + from kaiwu.core.gate import Gate + +logger = logging.getLogger(__name__) + +MAX_PARALLEL_WORKERS = 4 + + +class CycleError(Exception): + """Raised when the task DAG contains a cycle.""" + pass + + +class TaskCompiler: + """ + Lightweight DAG task scheduler. + Wraps PipelineOrchestrator.run() — each task in the DAG calls + orchestrator.run() with appropriate parameters. + """ + + def __init__( + self, + orchestrator: "PipelineOrchestrator", + gate: "Gate", + project_root: str, + ): + self.orchestrator = orchestrator + self.gate = gate + self.project_root = project_root + + def compile_and_run( + self, + tasks: list[dict], + on_status=None, + ) -> dict: + """ + Execute a DAG of tasks. + + Args: + tasks: list of task dicts, each with: + - "id": unique task identifier (str) + - "input": user_input string for the task + - "expert_type": (optional) gate classification override + - "depends_on": list of task IDs this task depends on + on_status: optional callback(stage, detail) + + Returns: + { + "results": {task_id: orchestrator_result_dict}, + "success": bool, # True if ALL tasks succeeded + "elapsed": float, + } + """ + start = time.time() + + if not tasks: + return {"results": {}, "success": True, "elapsed": 0.0} + + # Validate and build graph + task_map = {t["id"]: t for t in tasks} + self._validate_graph(task_map) + + # Topological layers (groups of tasks that can run in parallel) + layers = self._topological_layers(task_map) + + results: dict[str, dict] = {} + all_success = True + + for layer in layers: + if len(layer) == 1: + # Single task — run directly, no thread overhead + task_id = layer[0] + task_def = task_map[task_id] + result = self._execute_task(task_def, results, on_status) + results[task_id] = result + if not result["success"]: + all_success = False + else: + # Parallel execution + pool_size = min(len(layer), MAX_PARALLEL_WORKERS) + with ThreadPoolExecutor( + max_workers=pool_size, + thread_name_prefix="task_compiler", + ) as pool: + futures = {} + for task_id in layer: + task_def = task_map[task_id] + future = pool.submit( + self._execute_task, task_def, results, on_status + ) + futures[future] = task_id + + for future in as_completed(futures): + task_id = futures[future] + try: + result = future.result() + except Exception as e: + logger.error("Task %s raised: %s", task_id, e) + result = { + "success": False, + "context": None, + "error": str(e), + "elapsed": 0.0, + } + results[task_id] = result + if not result["success"]: + all_success = False + + elapsed = time.time() - start + return { + "results": results, + "success": all_success, + "elapsed": round(elapsed, 2), + } + + def _execute_task(self, task_def: dict, completed: dict, on_status) -> dict: + """Execute a single task via orchestrator.run().""" + task_id = task_def["id"] + user_input = task_def["input"] + + # Inject dependency context: append completed task outputs to input + deps = task_def.get("depends_on", []) + if deps: + dep_context = self._build_dependency_context(deps, completed) + if dep_context: + user_input = f"{user_input}\n\n[前置任务结果]\n{dep_context}" + + # Gate classification (use override or auto-classify) + expert_type = task_def.get("expert_type") + if expert_type: + gate_result = { + "expert_type": expert_type, + "task_summary": user_input[:20], + "difficulty": "easy", + } + else: + gate_result = self.gate.classify(user_input) + + logger.info("[task_compiler] Executing task %s: %s", task_id, user_input[:50]) + + return self.orchestrator.run( + user_input=user_input, + gate_result=gate_result, + project_root=self.project_root, + on_status=on_status, + ) + + @staticmethod + def _build_dependency_context(dep_ids: list[str], completed: dict) -> str: + """Build context string from completed dependency results.""" + parts = [] + for dep_id in dep_ids: + result = completed.get(dep_id) + if not result or not result.get("context"): + continue + ctx = result["context"] + # Extract explanation from generator output + gen = ctx.generator_output + if gen and gen.get("explanation"): + parts.append(f"任务{dep_id}: {gen['explanation'][:200]}") + elif gen and gen.get("patches"): + files = [p.get("file", "") for p in gen["patches"]] + parts.append(f"任务{dep_id}: 修改了 {', '.join(files)}") + return "\n".join(parts) + + @staticmethod + def _validate_graph(task_map: dict): + """Validate task graph: check for missing dependencies.""" + for task_id, task_def in task_map.items(): + for dep in task_def.get("depends_on", []): + if dep not in task_map: + raise ValueError( + f"Task '{task_id}' depends on '{dep}' which does not exist" + ) + + @staticmethod + def _topological_layers(task_map: dict) -> list[list[str]]: + """ + Kahn's algorithm producing layers of parallel-executable tasks. + Each layer contains tasks whose dependencies are all in previous layers. + Raises CycleError if the graph has a cycle. + """ + # Build adjacency and in-degree + in_degree = {tid: 0 for tid in task_map} + dependents = {tid: [] for tid in task_map} # tid -> list of tasks that depend on it + + for tid, task_def in task_map.items(): + for dep in task_def.get("depends_on", []): + in_degree[tid] += 1 + dependents[dep].append(tid) + + # Start with zero in-degree nodes + queue = deque(tid for tid, deg in in_degree.items() if deg == 0) + layers = [] + processed = 0 + + while queue: + # Current layer: all nodes with in_degree == 0 + layer = list(queue) + queue.clear() + layers.append(layer) + processed += len(layer) + + for tid in layer: + for dependent in dependents[tid]: + in_degree[dependent] -= 1 + if in_degree[dependent] == 0: + queue.append(dependent) + + if processed != len(task_map): + raise CycleError("Task DAG contains a cycle") + + return layers diff --git a/kaiwu/experts/search_augmentor.py b/kaiwu/experts/search_augmentor.py index b43b463..e824f27 100644 --- a/kaiwu/experts/search_augmentor.py +++ b/kaiwu/experts/search_augmentor.py @@ -110,9 +110,10 @@ class SearchAugmentorExpert: @staticmethod def _rerank_results(query: str, results: list[dict]) -> list[dict]: - """BM25 rerank: rescore search results by relevance to original query.""" + """BM25 rerank, then Cross-Encoder rerank if available (FLEX-2).""" if len(results) <= 1: return results + # Stage 1: BM25 rerank try: from rank_bm25 import BM25Plus corpus = [] @@ -124,9 +125,18 @@ class SearchAugmentorExpert: ranked = sorted( zip(results, scores), key=lambda x: x[1], reverse=True ) - return [r for r, _ in ranked] + results = [r for r, _ in ranked] except Exception: - return results + pass + + # Stage 2: Cross-Encoder rerank (optional, FLEX-2) + try: + from kaiwu.search.reranker import rerank + results = rerank(query, results, top_k=8) + except Exception: + pass + + return results def _extract(self, query: str, raw_results: str) -> str: """用LLM从原始搜索结果中提取关键信息。""" diff --git a/kaiwu/flywheel/ab_tester.py b/kaiwu/flywheel/ab_tester.py index 9b95a18..8fcf302 100644 --- a/kaiwu/flywheel/ab_tester.py +++ b/kaiwu/flywheel/ab_tester.py @@ -307,8 +307,13 @@ class ABTester: path = os.path.join(CANDIDATES_DIR, "candidates.json") data = {} for name, info in self._candidates.items(): + # Strip non-serializable fields from expert_def + expert_def = { + k: v for k, v in info["expert_def"].items() + if k != "_source" + } data[name] = { - "expert_def": info["expert_def"], + "expert_def": expert_def, "gate2_passed": info["gate2_passed"], "gate2_backtest": info.get("gate2_backtest", []), "backtest_success_rate": info.get("backtest_success_rate", 0.0), @@ -334,3 +339,28 @@ class ABTester: except Exception as e: logger.warning("Failed to load candidates: %s", e) self._candidates = {} + + # ── Prompt Optimization (SE-RED-4: 离线执行，使用外部API) ── + + def run_prompt_optimization( + self, + expert_name: str, + trajectories: list[TaskTrajectory], + api_key: str, + ) -> bool: + """ + AB测试通过后，分析成功轨迹优化YAML专家的system_prompt。 + SE-RED-4：使用外部API，用户需要提供API key。 + FLEX-1：无API key时跳过。 + """ + if not api_key: + logger.info("[ab_tester] 无API key，跳过prompt优化") + return False + + try: + from kaiwu.flywheel.prompt_optimizer import PromptOptimizer + optimizer = PromptOptimizer(api_key=api_key) + return optimizer.optimize_expert(expert_name, trajectories, self.registry) + except Exception as e: + logger.warning("[ab_tester] prompt优化失败: %s", e) + return False diff --git a/kaiwu/flywheel/prompt_optimizer.py b/kaiwu/flywheel/prompt_optimizer.py new file mode 100644 index 0000000..a2ef4f4 --- /dev/null +++ b/kaiwu/flywheel/prompt_optimizer.py @@ -0,0 +1,163 @@ +""" +Prompt Optimizer: analyzes successful trajectories and appends learned patterns +to the expert's YAML system_prompt field. +SE-RED-4: uses external API (Opus/Sonnet), offline execution. +""" + +import logging +import os +from typing import Optional + +import httpx +import yaml + +from kaiwu.flywheel.trajectory_collector import TaskTrajectory +from kaiwu.registry.expert_registry import ExpertRegistry + +logger = logging.getLogger(__name__) + +ANALYSIS_PROMPT = '''你是KWCode专家系统的prompt优化器。 + +分析以下{task_count}个成功任务的执行轨迹，提取可复用的经验规则。 + +## 轨迹摘要 +{trajectory_summary} + +## 当前专家system_prompt +``` +{current_prompt} +``` + +## 任务 +基于轨迹分析，生成2-5条具体的经验规则，格式如下： +- 每条规则一行，以"- "开头 +- 规则必须具体可操作（不要泛泛而谈） +- 规则应该帮助未来同类任务提高成功率 +- 不要重复已有prompt中的内容 + +只输出规则列表，不要解释。''' + + +class PromptOptimizer: + """ + Analyzes successful trajectories and appends learned patterns + to expert YAML system_prompt. + """ + + def __init__(self, api_key: str, model: str = "claude-sonnet-4-20250514"): + self.api_key = api_key + self.model = model + + def optimize_expert( + self, + expert_name: str, + trajectories: list[TaskTrajectory], + registry: ExpertRegistry, + ) -> bool: + """ + Analyze trajectories, generate conclusions, append to expert's system_prompt. + Returns True if optimization was applied. + """ + expert_def = registry.get(expert_name) + if not expert_def: + logger.warning("[prompt_optimizer] Expert not found: %s", expert_name) + return False + + source_path = expert_def.get("_source") + if not source_path or not os.path.isfile(source_path): + logger.warning("[prompt_optimizer] No source YAML for: %s", expert_name) + return False + + current_prompt = expert_def.get("system_prompt", "") + summary = self._summarize_trajectories(trajectories) + + # Call API for analysis + new_rules = self._call_api(len(trajectories), summary, current_prompt) + if not new_rules: + logger.info("[prompt_optimizer] API returned no rules") + return False + + # Append rules to system_prompt + updated_prompt = current_prompt.rstrip() + "\n\n## 经验规则（自动生成）\n" + new_rules + + # Write back to YAML + return self._update_yaml(source_path, updated_prompt, expert_name, registry) + + def _call_api(self, task_count: int, summary: str, current_prompt: str) -> Optional[str]: + """Call Opus/Sonnet API to generate optimization rules.""" + prompt = ANALYSIS_PROMPT.format( + task_count=task_count, + trajectory_summary=summary, + current_prompt=current_prompt[:2000], + ) + + try: + resp = httpx.post( + "https://api.anthropic.com/v1/messages", + headers={ + "x-api-key": self.api_key, + "anthropic-version": "2023-06-01", + "content-type": "application/json", + }, + json={ + "model": self.model, + "max_tokens": 1000, + "messages": [{"role": "user", "content": prompt}], + }, + timeout=60, + ) + resp.raise_for_status() + return resp.json()["content"][0]["text"].strip() + except Exception as e: + logger.error("[prompt_optimizer] API call failed: %s", e) + return None + + def _summarize_trajectories(self, trajectories: list[TaskTrajectory]) -> str: + """Extract patterns from trajectories for the API prompt.""" + from collections import Counter + + parts = [] + all_files = [] + for t in trajectories: + all_files.extend(t.files_modified) + + if all_files: + top_files = Counter(all_files).most_common(5) + parts.append("高频修改文件: " + ", ".join(f"{f}({c}次)" for f, c in top_files)) + + if trajectories: + avg_elapsed = sum(t.latency_s for t in trajectories) / len(trajectories) + parts.append(f"平均耗时: {avg_elapsed:.1f}秒") + avg_retries = sum(t.retry_count for t in trajectories) / len(trajectories) + parts.append(f"平均重试: {avg_retries:.1f}次") + + # Sample user inputs + inputs = [t.user_input[:80] for t in trajectories[:5]] + parts.append("典型任务: " + " | ".join(inputs)) + + return "\n".join(parts) if parts else "无足够数据" + + @staticmethod + def _update_yaml(source_path: str, new_prompt: str, expert_name: str, + registry: ExpertRegistry) -> bool: + """Write updated system_prompt back to YAML file.""" + try: + with open(source_path, "r", encoding="utf-8") as f: + data = yaml.safe_load(f) + + data["system_prompt"] = new_prompt + + with open(source_path, "w", encoding="utf-8") as f: + yaml.dump(data, f, default_flow_style=False, allow_unicode=True, + sort_keys=False, width=120) + + # Update in-memory registry + expert_def = registry.get(expert_name) + if expert_def: + expert_def["system_prompt"] = new_prompt + + logger.info("[prompt_optimizer] Updated system_prompt for %s", expert_name) + return True + except Exception as e: + logger.error("[prompt_optimizer] Failed to update YAML: %s", e) + return False diff --git a/kaiwu/memory/pattern_md.py b/kaiwu/memory/pattern_md.py index 55f544e..cc732a1 100644 --- a/kaiwu/memory/pattern_md.py +++ b/kaiwu/memory/pattern_md.py @@ -223,3 +223,123 @@ def show(project_root: str) -> str: return f.read() except Exception as e: return f"Failed to read PATTERN.md: {e}" + + +# ── Reflexion持久化 (SE-RED-5: 结构化格式) ── + +REFLECTION_TEMPLATE = "- [{date}] {task_summary} → {reflection}\n" + +_REFLECTION_FILE = "REFLECTION.md" + + +def _reflection_path(project_root: str) -> str: + return os.path.join(_kaiwu_dir(project_root), _REFLECTION_FILE) + + +def _read_reflection(project_root: str) -> str: + path = _reflection_path(project_root) + if not os.path.exists(path): + return "# KWCode Pattern Memory\n" + try: + with open(path, "r", encoding="utf-8") as f: + return f.read() + except Exception: + return "# KWCode Pattern Memory\n" + + +def _write_reflection(project_root: str, content: str): + _ensure_dir(project_root) + path = _reflection_path(project_root) + try: + with open(path, "w", encoding="utf-8") as f: + f.write(content) + except Exception as e: + logger.warning("Failed to write REFLECTION.md: %s", e) + + +def save_reflection( + project_root: str, + expert_type: str, + task_summary: str, + reflection: str, + success: bool, +): + """ + 把Reflection结果持久化到REFLECTION.md。 + SE-RED-5：结构化格式，不是自由文本。 + """ + from datetime import date as date_mod + + # 根据成败选择section + if success: + section = f"## {expert_type} 注意事项" + prefix = "注意" + else: + section = f"## {expert_type} 失败模式" + prefix = "根因" + + entry = REFLECTION_TEMPLATE.format( + date=date_mod.today().isoformat(), + task_summary=task_summary[:30], + reflection=f"{prefix}：{reflection[:80]}", + ) + + content = _read_reflection(project_root) + + # 找到对应section，追加条目 + if section in content: + content = content.replace( + section + "\n", + section + "\n" + entry, + ) + else: + content += f"\n{section}\n{entry}" + + # 限制每个section最多20条（防止无限增长） + content = _trim_reflection_sections(content, max_entries=20) + _write_reflection(project_root, content) + + +def get_reflections_for_plan(project_root: str, expert_type: str) -> str: + """ + /plan时读取相关历史Reflection，作为风险提示。 + 返回最近5条相关记录。 + """ + content = _read_reflection(project_root) + sections = [ + f"## {expert_type} 失败模式", + f"## {expert_type} 注意事项", + f"## {expert_type} 风险点", + ] + + result = [] + for section in sections: + if section in content: + after = content.split(section)[1].split("##")[0].strip() + recent = "\n".join(after.splitlines()[:5]) + if recent: + result.append(f"{section}\n{recent}") + + return "\n\n".join(result) + + +def _trim_reflection_sections(content: str, max_entries: int) -> str: + """每个section只保留最新的max_entries条。""" + lines = content.splitlines() + result = [] + current_section_entries = 0 + in_section = False + + for line in lines: + if line.startswith("## "): + in_section = True + current_section_entries = 0 + result.append(line) + elif in_section and line.startswith("- ["): + current_section_entries += 1 + if current_section_entries <= max_entries: + result.append(line) + else: + result.append(line) + + return "\n".join(result) diff --git a/kaiwu/search/reranker.py b/kaiwu/search/reranker.py new file mode 100644 index 0000000..8685941 --- /dev/null +++ b/kaiwu/search/reranker.py @@ -0,0 +1,88 @@ +""" +Cross-Encoder搜索结果重排。 +用CPU推理的轻量reranker，不需要GPU。 +FLEX-2：CPU性能低时跳过，只用BM25。 +""" + +import logging +import time +from typing import Optional + +logger = logging.getLogger(__name__) + +# 使用sentence-transformers的cross-encoder +# 模型：cross-encoder/ms-marco-MiniLM-L-6-v2（82MB，CPU可跑） +RERANKER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2" +_reranker = None +_reranker_disabled = False # 如果加载失败或太慢，后续跳过 + + +def get_reranker(): + """懒加载，首次使用时下载模型。""" + global _reranker, _reranker_disabled + if _reranker_disabled: + return None + if _reranker is None: + try: + from sentence_transformers import CrossEncoder + _reranker = CrossEncoder(RERANKER_MODEL) + logger.info("[reranker] 模型加载完成: %s", RERANKER_MODEL) + except ImportError: + logger.info("[reranker] sentence-transformers未安装，跳过重排") + _reranker_disabled = True + except Exception as e: + logger.warning("[reranker] 模型加载失败: %s，跳过重排", e) + _reranker_disabled = True + return _reranker + + +def rerank( + query: str, + results: list[dict], + top_k: int = 3, +) -> list[dict]: + """ + 用Cross-Encoder对搜索结果重排。 + 失败时降级返回原始顺序（FLEX-2）。 + + results: [{"title": str, "url": str, "snippet": str}, ...] + """ + global _reranker_disabled + + reranker = get_reranker() + if not reranker or not results: + return results[:top_k] + + # 构建(query, document)对 + pairs = [ + (query, f"{r.get('title', '')} {r.get('snippet', '')}") + for r in results + ] + + try: + t0 = time.perf_counter() + scores = reranker.predict(pairs) + elapsed_ms = (time.perf_counter() - t0) * 1000 + + # 按分数排序 + ranked = sorted( + zip(scores, results), + key=lambda x: x[0], + reverse=True, + ) + + logger.info( + "[reranker] 重排完成：%d→%d结果，耗时%.0fms", + len(results), top_k, elapsed_ms, + ) + + # FLEX-2：耗时超过2秒说明CPU太慢，后续跳过 + if elapsed_ms > 2000: + logger.warning("[reranker] 耗时%.0fms超过2s，后续跳过重排", elapsed_ms) + _reranker_disabled = True + + return [r for _, r in ranked[:top_k]] + + except Exception as e: + logger.warning("[reranker] 重排失败: %s，返回原始顺序", e) + return results[:top_k] diff --git a/kaiwu/tests/test_task_compiler.py b/kaiwu/tests/test_task_compiler.py new file mode 100644 index 0000000..f7bd605 --- /dev/null +++ b/kaiwu/tests/test_task_compiler.py @@ -0,0 +1,273 @@ +""" +Tests for TaskCompiler: lightweight DAG task scheduler. +Tests serial (dependency chain) and parallel (independent tasks) scenarios. +Uses mock orchestrator since Ollama may not be running. +""" + +import time +from unittest.mock import MagicMock, patch + +import pytest + +from kaiwu.core.task_compiler import TaskCompiler, CycleError + + +# ── Fixtures ── + +def _make_mock_orchestrator(): + """Create a mock orchestrator that returns success with realistic context.""" + orch = MagicMock() + + def mock_run(user_input, gate_result, project_root, on_status=None, no_search=False): + # Simulate some work + time.sleep(0.05) + ctx = MagicMock() + ctx.generator_output = { + "patches": [{"file": "src/main.py", "original": "old", "modified": "new"}], + "explanation": f"Completed: {user_input[:30]}", + } + ctx.user_input = user_input + return { + "success": True, + "context": ctx, + "error": None, + "elapsed": 0.05, + } + + orch.run = MagicMock(side_effect=mock_run) + return orch + + +def _make_mock_gate(): + """Create a mock gate that classifies everything as codegen.""" + gate = MagicMock() + gate.classify = MagicMock(return_value={ + "expert_type": "codegen", + "task_summary": "test task", + "difficulty": "easy", + }) + return gate + + +# ── Serial Tests ── + +class TestTaskCompilerSerial: + """Serial scenario: refactor → write tests (t2 depends on t1).""" + + def test_serial_execution_order(self): + """t2 must execute after t1 completes.""" + orch = _make_mock_orchestrator() + gate = _make_mock_gate() + compiler = TaskCompiler(orchestrator=orch, gate=gate, project_root="/tmp/test") + + tasks = [ + {"id": "t1", "input": "refactor extract_data into two functions", "depends_on": []}, + {"id": "t2", "input": "write tests for the new functions", "depends_on": ["t1"]}, + ] + + result = compiler.compile_and_run(tasks) + + assert result["success"] is True + assert "t1" in result["results"] + assert "t2" in result["results"] + assert result["results"]["t1"]["success"] is True + assert result["results"]["t2"]["success"] is True + + # Verify t1 was called before t2 (check call order) + calls = orch.run.call_args_list + assert len(calls) == 2 + # First call should be t1's input + assert "refactor" in calls[0].kwargs.get("user_input", calls[0][1]["user_input"] if len(calls[0]) > 1 else calls[0][0][0]) + + def test_serial_context_injection(self): + """t2 should receive t1's output in its input.""" + orch = _make_mock_orchestrator() + gate = _make_mock_gate() + compiler = TaskCompiler(orchestrator=orch, gate=gate, project_root="/tmp/test") + + tasks = [ + {"id": "t1", "input": "refactor extract_data", "depends_on": []}, + {"id": "t2", "input": "write tests", "depends_on": ["t1"]}, + ] + + result = compiler.compile_and_run(tasks) + assert result["success"] is True + + # t2's user_input should contain dependency context + calls = orch.run.call_args_list + t2_input = calls[1][1]["user_input"] if "user_input" in (calls[1][1] if len(calls[1]) > 1 else {}) else calls[1].kwargs.get("user_input", "") + assert "前置任务结果" in t2_input + + def test_serial_failure_propagation(self): + """If t1 fails, t2 still runs but without dependency context.""" + orch = MagicMock() + call_count = [0] + + def mock_run(**kwargs): + call_count[0] += 1 + if call_count[0] == 1: + # t1 fails + ctx = MagicMock() + ctx.generator_output = None + return {"success": False, "context": ctx, "error": "failed", "elapsed": 0.1} + else: + # t2 succeeds + ctx = MagicMock() + ctx.generator_output = {"patches": [], "explanation": "done"} + return {"success": True, "context": ctx, "error": None, "elapsed": 0.1} + + orch.run = MagicMock(side_effect=mock_run) + gate = _make_mock_gate() + compiler = TaskCompiler(orchestrator=orch, gate=gate, project_root="/tmp/test") + + tasks = [ + {"id": "t1", "input": "task1", "depends_on": []}, + {"id": "t2", "input": "task2", "depends_on": ["t1"]}, + ] + + result = compiler.compile_and_run(tasks) + # Overall should be False because t1 failed + assert result["success"] is False + assert result["results"]["t1"]["success"] is False + assert result["results"]["t2"]["success"] is True + + +# ── Parallel Tests ── + +class TestTaskCompilerParallel: + """Parallel scenario: 3 independent tasks run concurrently.""" + + def test_parallel_all_succeed(self): + """Three independent tasks should all succeed.""" + orch = _make_mock_orchestrator() + gate = _make_mock_gate() + compiler = TaskCompiler(orchestrator=orch, gate=gate, project_root="/tmp/test") + + tasks = [ + {"id": "t1", "input": "add comments to function_a", "depends_on": []}, + {"id": "t2", "input": "add comments to function_b", "depends_on": []}, + {"id": "t3", "input": "add comments to function_c", "depends_on": []}, + ] + + result = compiler.compile_and_run(tasks) + + assert result["success"] is True + assert len(result["results"]) == 3 + for tid in ["t1", "t2", "t3"]: + assert result["results"][tid]["success"] is True + + def test_parallel_faster_than_serial(self): + """Parallel execution should be faster than serial (3 * 50ms > total).""" + orch = _make_mock_orchestrator() + gate = _make_mock_gate() + compiler = TaskCompiler(orchestrator=orch, gate=gate, project_root="/tmp/test") + + tasks = [ + {"id": "t1", "input": "task a", "depends_on": []}, + {"id": "t2", "input": "task b", "depends_on": []}, + {"id": "t3", "input": "task c", "depends_on": []}, + ] + + result = compiler.compile_and_run(tasks) + + # Each task takes ~50ms. Serial would be ~150ms. Parallel should be ~50-80ms. + assert result["elapsed"] < 0.15, f"Parallel took {result['elapsed']}s, expected < 0.15s" + + def test_parallel_with_expert_type_override(self): + """Tasks with explicit expert_type should skip gate classification.""" + orch = _make_mock_orchestrator() + gate = _make_mock_gate() + compiler = TaskCompiler(orchestrator=orch, gate=gate, project_root="/tmp/test") + + tasks = [ + {"id": "t1", "input": "add docstring", "expert_type": "doc", "depends_on": []}, + {"id": "t2", "input": "fix bug", "expert_type": "locator_repair", "depends_on": []}, + ] + + result = compiler.compile_and_run(tasks) + assert result["success"] is True + + # Gate should NOT have been called (expert_type was pre-specified) + gate.classify.assert_not_called() + + +# ── Validation Tests ── + +class TestTaskCompilerValidation: + """Edge cases and error handling.""" + + def test_empty_task_list(self): + """Empty task list should return immediately.""" + orch = _make_mock_orchestrator() + gate = _make_mock_gate() + compiler = TaskCompiler(orchestrator=orch, gate=gate, project_root="/tmp/test") + + result = compiler.compile_and_run([]) + assert result["success"] is True + assert result["results"] == {} + + def test_single_task(self): + """Single task with no dependencies.""" + orch = _make_mock_orchestrator() + gate = _make_mock_gate() + compiler = TaskCompiler(orchestrator=orch, gate=gate, project_root="/tmp/test") + + tasks = [{"id": "only", "input": "do something", "depends_on": []}] + result = compiler.compile_and_run(tasks) + assert result["success"] is True + assert "only" in result["results"] + + def test_missing_dependency_raises(self): + """Referencing a non-existent dependency should raise ValueError.""" + orch = _make_mock_orchestrator() + gate = _make_mock_gate() + compiler = TaskCompiler(orchestrator=orch, gate=gate, project_root="/tmp/test") + + tasks = [ + {"id": "t1", "input": "task", "depends_on": ["nonexistent"]}, + ] + + with pytest.raises(ValueError, match="does not exist"): + compiler.compile_and_run(tasks) + + def test_cycle_detection(self): + """Circular dependencies should raise CycleError.""" + orch = _make_mock_orchestrator() + gate = _make_mock_gate() + compiler = TaskCompiler(orchestrator=orch, gate=gate, project_root="/tmp/test") + + tasks = [ + {"id": "t1", "input": "task1", "depends_on": ["t2"]}, + {"id": "t2", "input": "task2", "depends_on": ["t1"]}, + ] + + with pytest.raises(CycleError): + compiler.compile_and_run(tasks) + + def test_gate_auto_classify(self): + """Tasks without expert_type should use gate.classify().""" + orch = _make_mock_orchestrator() + gate = _make_mock_gate() + compiler = TaskCompiler(orchestrator=orch, gate=gate, project_root="/tmp/test") + + tasks = [{"id": "t1", "input": "fix the bug in login", "depends_on": []}] + result = compiler.compile_and_run(tasks) + assert result["success"] is True + gate.classify.assert_called_once() + + def test_diamond_dependency(self): + """Diamond DAG: t1 → t2, t1 → t3, t2+t3 → t4.""" + orch = _make_mock_orchestrator() + gate = _make_mock_gate() + compiler = TaskCompiler(orchestrator=orch, gate=gate, project_root="/tmp/test") + + tasks = [ + {"id": "t1", "input": "setup", "depends_on": []}, + {"id": "t2", "input": "branch a", "depends_on": ["t1"]}, + {"id": "t3", "input": "branch b", "depends_on": ["t1"]}, + {"id": "t4", "input": "merge", "depends_on": ["t2", "t3"]}, + ] + + result = compiler.compile_and_run(tasks) + assert result["success"] is True + assert len(result["results"]) == 4 diff --git a/pyproject.toml b/pyproject.toml index aeef691..667456d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "kwcode" -version = "0.7.0" +version = "0.8.0" description = "KwCode - Local-model coding agent with MoE expert pipeline" requires-python = ">=3.10" @@ -28,6 +28,9 @@ dependencies = [ "aiosqlite>=0.20.0", ] +[project.optional-dependencies] +rerank = ["sentence-transformers>=2.7.0"] + [project.scripts] kwcode = "kaiwu.cli.main:app"