Add CodePilot CLI tool with feedback collection

codepilot.py

@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+CodePilot — 你的專屬 AI 程式助手
+=================================
+
+像 Claude Code 的終端 CLI 工具，內建：
+  🤖 本地 Qwen2.5-Coder 模型推理
+  📁 讀取/修改你的專案文件
+  👍👎 一鍵回饋，自動收集訓練數據
+  🔄 定期用你的回饋數據訓練，模型越用越聰明
+
+Install:
+    pip install transformers peft bitsandbytes accelerate trl datasets rich
+
+Usage:
+    python codepilot.py                          # 啟動
+    python codepilot.py --project ~/my-project   # 指定專案
+    python codepilot.py --adapter ./my-adapter   # 用微調模型
+    python codepilot.py --stats                  # 回饋統計
+    python codepilot.py --train                  # 訓練模型
+"""
+import argparse, json, os, sqlite3, subprocess, sys, textwrap, torch
+from datetime import datetime
+from pathlib import Path
+
+DEFAULT_MODEL = "Qwen/Qwen2.5-Coder-3B-Instruct"
+CONFIG_DIR = os.path.expanduser("~/.codepilot")
+DB_PATH = os.path.join(CONFIG_DIR, "feedback.db")
+AUTO_TRAIN_THRESHOLD = 50
+
+class FeedbackDB:
+    def __init__(self, db_path=DB_PATH):
+        os.makedirs(os.path.dirname(db_path), exist_ok=True)
+        self.conn = sqlite3.connect(db_path)
+        self.conn.execute("""CREATE TABLE IF NOT EXISTS feedback (
+            id INTEGER PRIMARY KEY, timestamp TEXT, prompt TEXT, completion TEXT,
+            label INTEGER, edited_completion TEXT, project_dir TEXT, files_context TEXT)""")
+        self.conn.commit()
+
+    def save(self, prompt, completion, label, edited=None, project_dir=None, files=None):
+        self.conn.execute("INSERT INTO feedback VALUES (NULL,?,?,?,?,?,?,?)",
+            (datetime.now().isoformat(), prompt, completion, int(label), edited, project_dir,
+             json.dumps(files) if files else None))
+        self.conn.commit()
+
+    def count(self):
+        r = self.conn.execute("SELECT COUNT(*), SUM(label), SUM(CASE WHEN edited_completion IS NOT NULL THEN 1 ELSE 0 END) FROM feedback").fetchone()
+        return {"total": r[0] or 0, "thumbs_up": r[1] or 0, "edits": r[2] or 0}
+
+    def export_kto(self):
+        rows = self.conn.execute("SELECT prompt, completion, label FROM feedback").fetchall()
+        return [{"prompt":[{"role":"user","content":p}],"completion":[{"role":"assistant","content":c}],"label":bool(l)} for p,c,l in rows]
+
+    def export_sft(self):
+        rows = self.conn.execute("SELECT prompt, edited_completion FROM feedback WHERE edited_completion IS NOT NULL").fetchall()
+        return [{"messages":[{"role":"user","content":p},{"role":"assistant","content":c}]} for p,c in rows]
+
+    def export_dpo(self):
+        rows = self.conn.execute("SELECT prompt, completion, edited_completion FROM feedback WHERE edited_completion IS NOT NULL").fetchall()
+        return [{"prompt":[{"role":"user","content":p}],"chosen":[{"role":"assistant","content":e}],"rejected":[{"role":"assistant","content":o}]} for p,o,e in rows]
+
+class CodeModel:
+    def __init__(self, model_name=DEFAULT_MODEL, adapter_path=None):
+        from transformers import AutoTokenizer, AutoModelForCausalLM
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
+        if adapter_path and os.path.exists(adapter_path):
+            from peft import PeftModel
+            self.model = PeftModel.from_pretrained(self.model, adapter_path)
+        self.model.eval()
+
+    def generate(self, user_message, system_prompt=None, file_context=None, max_tokens=2048):
+        messages = []
+        if system_prompt: messages.append({"role":"system","content":system_prompt})
+        if file_context:
+            messages.append({"role":"user","content":f"相關文件:\n\n{file_context}"})
+            messages.append({"role":"assistant","content":"已了解。請問需要什麼幫助？"})
+        messages.append({"role":"user","content":user_message})
+        text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)
+        with torch.no_grad():
+            outputs = self.model.generate(**inputs, max_new_tokens=max_tokens, do_sample=True, temperature=0.7, top_p=0.9, repetition_penalty=1.1, pad_token_id=self.tokenizer.pad_token_id)
+        return self.tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+
+def extract_code_blocks(text):
+    blocks = []
+    parts = text.split("```")
+    for i in range(1, len(parts), 2):
+        lines = parts[i].split("\n", 1)
+        lang = lines[0].strip() or "python"
+        code = lines[1].strip() if len(lines) > 1 else ""
+        blocks.append((lang, code))
+    return blocks
+
+def run_cli(args):
+    from rich.console import Console
+    from rich.markdown import Markdown
+    from rich.panel import Panel
+    from rich.prompt import Prompt
+    from rich.syntax import Syntax
+    from rich.table import Table
+
+    console = Console(); db = FeedbackDB()
+    console.print(Panel.fit("[bold cyan]CodePilot[/] — 你的專屬 AI 程式助手\n" + f"[dim]Model: {args.model or DEFAULT_MODEL}[/]", border_style="cyan"))
+
+    with console.status("[bold green]載入模型中..."):
+        model = CodeModel(args.model or DEFAULT_MODEL, args.adapter)
+    console.print("[green]✅ 模型載入完成[/]\n")
+    console.print("[dim]指令: /file <path> 讀文件 | /apply 套用code | /stats 統計 | /train 訓練 | /quit 退出[/]\n")
+
+    system_prompt = "You are CodePilot, an expert programming assistant. Write clean, efficient, well-documented code. When modifying existing code, show the complete modified version."
+    project_dir = args.project or os.getcwd()
+    file_context = current_response = current_prompt = None
+
+    while True:
+        try: user_input = Prompt.ask("\n[bold green]🧑 You")
+        except (EOFError, KeyboardInterrupt): break
+        if not user_input.strip(): continue
+        cmd = user_input.strip().lower()
+
+        if cmd in ("/quit", "/exit"): break
+        elif cmd == "/stats":
+            stats = db.count()
+            t = Table(title="📊 回饋統計"); t.add_column("指標",style="cyan"); t.add_column("數值",style="green")
+            t.add_row("總回饋",str(stats["total"])); t.add_row("👍",str(stats["thumbs_up"]))
+            t.add_row("👎",str(stats["total"]-stats["thumbs_up"])); t.add_row("✏️修改",str(stats["edits"]))
+            console.print(t); continue
+        elif cmd == "/train": trigger_cli_training(db, console, args); continue
+        elif cmd.startswith("/file "):
+            fp = os.path.join(project_dir, user_input[6:].strip())
+            if os.path.exists(fp):
+                with open(fp) as f: content = f.read()
+                file_context = f"--- {fp} ---\n{content}\n--- END ---"
+                console.print(f"[green]📁 已讀取: {fp} ({len(content)} chars)[/]")
+            else: console.print(f"[red]❌ 不存在: {fp}[/]")
+            continue
+        elif cmd == "/apply":
+            if current_response:
+                for i,(lang,code) in enumerate(extract_code_blocks(current_response)):
+                    console.print(Syntax(code, lang or "python", theme="monokai"))
+                    fp = Prompt.ask("  儲存到? (Enter跳過)")
+                    if fp.strip():
+                        full = os.path.join(project_dir, fp)
+                        os.makedirs(os.path.dirname(full) or ".", exist_ok=True)
+                        open(full,"w").write(code)
+                        console.print(f"  [green]✅ {full}[/]")
+            continue
+
+        current_prompt = user_input
+        with console.status("[bold cyan]思考中..."):
+            current_response = model.generate(user_input, system_prompt=system_prompt, file_context=file_context)
+        console.print("\n[bold blue]🤖 CodePilot:[/]")
+        console.print(Markdown(current_response))
+        console.print("\n[dim][green]y[/]=👍  [red]n[/]=👎  [yellow]e[/]=✏️修改  Enter=跳過[/]")
+        fb = Prompt.ask("  ", choices=["y","n","e",""], default="", show_choices=False)
+
+        if fb == "y":
+            db.save(current_prompt, current_response, label=1, project_dir=project_dir)
+            s = db.count(); console.print(f"  [green]👍 +1 (累計:{s['total']})[/]")
+        elif fb == "n":
+            db.save(current_prompt, current_response, label=0, project_dir=project_dir)
+            s = db.count(); console.print(f"  [red]👎 +1 (累計:{s['total']})[/]")
+        elif fb == "e":
+            console.print("  [yellow]貼上修改版 (END 結束):[/]")
+            lines = []
+            while True:
+                try:
+                    l = input()
+                    if l.strip()=="END": break
+                    lines.append(l)
+                except EOFError: break
+            edited = "\n".join(lines)
+            if edited.strip():
+                db.save(current_prompt, current_response, label=1, edited=edited, project_dir=project_dir)
+                s = db.count(); console.print(f"  [yellow]✏️ +1 (累計:{s['total']}, 修改:{s['edits']})[/]")
+
+        if db.count()["total"] % AUTO_TRAIN_THRESHOLD == 0 and db.count()["total"] > 0:
+            console.print(f"\n  [bold yellow]🔔 累積 {db.count()['total']} 條！codepilot --train[/]")
+
+    console.print("\n[cyan]👋 再見！[/]")
+
+def trigger_cli_training(db, console, args):
+    stats = db.count()
+    if stats["total"] == 0: console.print("[yellow]⚠️ 無數據[/]"); return
+    console.print(f"\n[bold]🚀 訓練[/] 👍:{stats['thumbs_up']} 👎:{stats['total']-stats['thumbs_up']} ✏️:{stats['edits']}")
+    from datasets import Dataset
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+    from peft import LoraConfig, prepare_model_for_kbit_training
+    model_name = args.model or DEFAULT_MODEL
+    output_dir = os.path.join(CONFIG_DIR, f"adapter_{datetime.now().strftime('%Y%m%d_%H%M')}")
+    bnb = BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type="nf4",bnb_4bit_compute_dtype=torch.bfloat16,bnb_4bit_use_double_quant=True)
+    peft_cfg = LoraConfig(r=16,lora_alpha=32,lora_dropout=0.05,bias="none",task_type="CAUSAL_LM",target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"])
+
+    sft_data = db.export_sft()
+    kto_data = db.export_kto()
+
+    if sft_data:
+        console.print(f"\n[bold]📚 SFT ({len(sft_data)} edits)...[/]")
+        from trl import SFTTrainer, SFTConfig
+        ds = Dataset.from_list(sft_data)
+        model = AutoModelForCausalLM.from_pretrained(model_name,quantization_config=bnb,device_map="auto",trust_remote_code=True)
+        tok = AutoTokenizer.from_pretrained(model_name)
+        if tok.pad_token is None: tok.pad_token = tok.eos_token
+        model = prepare_model_for_kbit_training(model)
+        trainer = SFTTrainer(model=model,args=SFTConfig(output_dir=output_dir,learning_rate=2e-4,num_train_epochs=3,per_device_train_batch_size=1,gradient_accumulation_steps=8,max_seq_length=1024,gradient_checkpointing=True,bf16=True,optim="paged_adamw_8bit",logging_steps=5,save_total_limit=1,logging_strategy="steps",logging_first_step=True),processing_class=tok,train_dataset=ds,peft_config=peft_cfg)
+        trainer.train(); trainer.save_model(output_dir)
+        console.print(f"[green]✅ SFT 完成[/]"); del model; torch.cuda.empty_cache()
+    elif len(kto_data) >= 10:
+        console.print(f"\n[bold]📚 KTO ({len(kto_data)} feedbacks)...[/]")
+        from trl import KTOConfig, KTOTrainer
+        ds = Dataset.from_list(kto_data)
+        model = AutoModelForCausalLM.from_pretrained(model_name,quantization_config=bnb,device_map="auto",trust_remote_code=True)
+        tok = AutoTokenizer.from_pretrained(model_name)
+        if tok.pad_token is None: tok.pad_token = tok.eos_token
+        trainer = KTOTrainer(model=model,args=KTOConfig(output_dir=output_dir,learning_rate=1e-5,num_train_epochs=1,per_device_train_batch_size=1,gradient_accumulation_steps=8,max_length=1024,gradient_checkpointing=True,bf16=True,logging_steps=5,logging_strategy="steps",logging_first_step=True),processing_class=tok,train_dataset=ds,peft_config=peft_cfg)
+        trainer.train(); trainer.save_model(output_dir)
+        console.print(f"[green]✅ KTO 完成[/]")
+
+    console.print(f"\n[bold green]🎉 訓練完成！[/]\n   Adapter: {output_dir}\n   重啟: codepilot --adapter {output_dir}")
+
+def show_stats():
+    from rich.console import Console; from rich.table import Table
+    console = Console(); db = FeedbackDB(); s = db.count()
+    t = Table(title="📊 CodePilot 回饋統計"); t.add_column("指標",style="cyan"); t.add_column("數值",style="green"); t.add_column("",style="dim")
+    t.add_row("總回饋",str(s["total"]),"█"*min(s["total"]//2,40))
+    t.add_row("👍",str(s["thumbs_up"]),"█"*min(s["thumbs_up"]//2,40))
+    t.add_row("👎",str(s["total"]-s["thumbs_up"]),"█"*min((s["total"]-s["thumbs_up"])//2,40))
+    t.add_row("✏️修改",str(s["edits"]),"█"*min(s["edits"]//2,40))
+    console.print(t)
+    if s["total"]>0:
+        r = s["thumbs_up"]/s["total"]*100; console.print(f"\n   接受率: {r:.0f}%")
+        if r < 50: console.print("   [yellow]💡 接受率低，建議 --train[/]")
+
+def main():
+    p = argparse.ArgumentParser(description="CodePilot — 你的專屬 AI 程式助手")
+    p.add_argument("--model",type=str,default=None,help=f"模型 (預設:{DEFAULT_MODEL})")
+    p.add_argument("--adapter",type=str,default=None,help="LoRA adapter")
+    p.add_argument("--project",type=str,default=None,help="專案目錄")
+    p.add_argument("--stats",action="store_true",help="回饋統計")
+    p.add_argument("--train",action="store_true",help="訓練模型")
+    a = p.parse_args()
+    if a.stats: show_stats()
+    elif a.train: from rich.console import Console; trigger_cli_training(FeedbackDB(), Console(), a)
+    else: run_cli(a)
+
+if __name__ == "__main__": main()