codepilot.py

#!/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()