Upload second_brain.py with huggingface_hub

second_brain.py

@@ -0,0 +1,440 @@
+#!/usr/bin/env python3
+"""
+FSI_FELON · SECOND BRAIN LOOP RUNNER
+Complete closed-loop: generate → verify (Veritas) → tool (nanobots) → fix → repeat
+This IS the SOTA differentiator — architecture, not parameter count.
+"""
+
+import os, sys, json, time, re, subprocess, tempfile
+from pathlib import Path
+from dataclasses import dataclass, field
+from typing import Optional, List, Dict, Any
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from nanobot_tools import SwarmCoordinator
+from quantum.gated_conv_engine import FelonGatedConvModel
+from quantum.bpe_tokenizer import BPETokenizer
+from tool_schema import tool_signature, TOOL_SCHEMA
+
+TOKENIZER_PATH = "quantum/felon_bpe_tokenizer.json"
+MODEL_PATH = "felon_gatedconv_best.pt"  # will be felon_fast_best.pt when ready
+WORK_DIR = "/tmp/fsi_second_brain"
+MAX_ITERATIONS = 12
+
+os.makedirs(WORK_DIR, exist_ok=True)
+
+
+@dataclass
+class ToolResult:
+    success: bool
+    result: str = ""
+    error: str = ""
+
+@dataclass
+class Step:
+    thought: str
+    tool_call: Optional[Dict] = None
+    result: Optional[ToolResult] = None
+    iteration: int = 0
+
+@dataclass
+class VeritasResult:
+    passed: bool
+    score: int  # 0-6
+    details: Dict[str, bool]
+    errors: List[str]
+
+
+# ═══════════════════════════════════════════════════════════
+# VERITAS LAYER — 6 hard checks
+# ═══════════════════════════════════════════════════════════
+
+class VeritasLayer:
+    """6-check verifier. Returns score 0-6. >=5 = pass."""
+    
+    SECURITY_PATTERNS = [
+        r"eval\(", r"exec\(", r"__import__", r"subprocess\.call",
+        r"os\.system", r"pickle\.loads", r"base64\.b64decode",
+    ]
+    
+    def __init__(self):
+        self.checks = ["syntax", "imports", "security", "logic", "tests", "edge_cases"]
+    
+    def run(self, code: str, filepath: str = "") -> VeritasResult:
+        details = {}
+        errors = []
+        
+        # 1. Syntax
+        try:
+            compile(code, filepath or "<veritas>", "exec")
+            details["syntax"] = True
+        except SyntaxError as e:
+            details["syntax"] = False
+            errors.append(f"SyntaxError: {e}")
+        
+        # 2. Imports
+        import re
+        imports = re.findall(r"^(?:from|import)\s+(\w+)", code, re.MULTILINE)
+        ok = True
+        for mod in imports:
+            try:
+                __import__(mod)
+            except ImportError:
+                ok = False
+                errors.append(f"Missing import: {mod}")
+        details["imports"] = ok
+        
+        # 3. Security
+        secure = True
+        for pattern in self.SECURITY_PATTERNS:
+            if re.search(pattern, code):
+                secure = False
+                errors.append(f"Security: found {pattern}")
+        details["security"] = secure
+        
+        # 4. Logic (has function/class/control flow)
+        has_logic = any(kw in code for kw in ["def ", "class ", "return ", "for ", "while ", "if ", "lambda", "import"])
+        details["logic"] = has_logic
+        
+        # 5. Tests (has assert)
+        has_tests = "assert" in code
+        details["tests"] = has_tests
+        
+        # 6. Edge cases (has try/except or if/else)
+        has_edges = ("try:" in code and "except" in code) or ("if " in code and "else" in code)
+        details["edge_cases"] = has_edges
+        
+        score = sum(details.values())
+        passed = score >= 5
+        
+        return VeritasResult(passed, score, details, errors)
+
+
+# ═══════════════════════════════════════════════════════════
+# NANOBOT TOOL EXECUTORS
+# ═══════════════════════════════════════════════════════════
+
+class ToolExecutor:
+    """Wraps nanobot swarm with clean result handling."""
+    
+    def __init__(self):
+        self.swarm = SwarmCoordinator()
+    
+    def execute(self, tool_call: Dict) -> ToolResult:
+        if not tool_call:
+            return ToolResult(False, error="Empty tool call")
+        
+        if tool_call.get("tool") == "done":
+            return ToolResult(True, "TASK_COMPLETE")
+        
+        tool_name = tool_call.get("tool")
+        action = tool_call.get("action", "default")
+        args = {k: v for k, v in tool_call.items() if k not in ("tool", "action")}
+        
+        # Special handling for fs action parameter
+        if tool_name == "fs" and "action" in tool_call:
+            args["action"] = tool_call["action"]
+        
+        try:
+            success, result, error = self.swarm.dispatch(tool_name, args)
+            return ToolResult(success, str(result)[:1000] if result else "", error or "")
+        except Exception as e:
+            return ToolResult(False, error=str(e))
+
+
+# ═══════════════════════════════════════════════════════════
+# SECOND BRAIN ORCHESTRATOR
+# ═══════════════════════════════════════════════════════════
+
+class SecondBrain:
+    def __init__(self, model: Optional[FelonGatedConvModel] = None, 
+                 tokenizer: Optional[BPETokenizer] = None):
+        self.model = model
+        self.tok = tokenizer
+        self.executor = ToolExecutor()
+        self.veritas = VeritasLayer()
+        self.history: List[Step] = []
+        os.makedirs(WORK_DIR, exist_ok=True)
+    
+    def load_model(self):
+        if self.model is not None:
+            return
+        self.tok = BPETokenizer(vocab_size=4096)
+        self.tok.load(TOKENIZER_PATH)
+        self.model = FelonGatedConvModel(
+            vocab_size=self.tok.vocab_size,
+            hidden=256, n_layers=10, n_heads=4, n_kv_heads=2,
+            inter=512, kernel=3, max_seq=512, dropout=0.1
+        )
+        if os.path.exists(MODEL_PATH):
+            self.model.load(MODEL_PATH)
+            print(f"  Loaded model: {MODEL_PATH}")
+        else:
+            print(f"  WARNING: No model at {MODEL_PATH} — using random weights")
+    
+    def build_prompt(self, task: str) -> str:
+        # Match training data format exactly — NO tool signature
+        return f"TASK: {task}\nTHOUGHT: "
+
+    def add_to_history(self, step: Step):
+        self.history.append(step)
+
+    def get_context(self, max_chars: int = 3000) -> str:
+        ctx = ""
+        for i, step in enumerate(self.history):
+            if i == 0 and step.thought:
+                ctx += step.thought + ("\n" if not step.thought.endswith("\n") else "")
+                continue
+            if step.thought:
+                ctx += f"THOUGHT: {step.thought}\n"
+            if step.tool_call:
+                ctx += f"{json.dumps(step.tool_call)}\n"
+            if step.result:
+                if step.result.success:
+                    ctx += f"RESULT: {step.result.result}\n"
+                else:
+                    ctx += f"ERROR: {step.result.error}\n"
+        return ctx[-max_chars:]
+    
+    def generate_next(self, max_new: int = 200) -> str:
+        ids = self.tok.encode(self.get_context())
+        if len(ids) > self.model.max_seq - max_new:
+            ids = ids[-(self.model.max_seq - max_new):]
+        out = self.model.generate(ids, max_new=max_new, temp=0.7, top_k=50, eos_id=self.tok.EOS)
+        return self.tok.decode(out)
+        if len(ids) > self.model.max_seq - max_new:
+            ids = ids[-(self.model.max_seq - max_new):]
+        out = self.model.generate(ids, max_new=max_new, temp=0.7, top_k=50, eos_id=self.tok.EOS)
+        return self.tok.decode(out)
+    
+    def parse_tool_call(self, text: str) -> Optional[Dict]:
+        matches = re.findall(r'\{[^}]*"tool"[^}]*\}', text)
+        for m in matches:
+            try:
+                return json.loads(m)
+            except json.JSONDecodeError:
+                continue
+        return None
+    
+    def extract_thought(self, text: str, tool_call: Optional[Dict]) -> str:
+        if not tool_call:
+            return text.strip()
+        tc_str = json.dumps(tool_call)
+        idx = text.find(tc_str)
+        if idx > 0:
+            thought = text[:idx].strip()
+            # Remove "THOUGHT: " prefix if present
+            if thought.startswith("THOUGHT: "):
+                thought = thought[9:]
+            return thought
+        return text.strip()
+    
+    def write_code_files(self, code: str, base_name: str = "main") -> List[str]:
+        """Extract and write Python files from code blocks."""
+        written = []
+        # Find ```python ... ``` blocks
+        blocks = re.findall(r'```python\n(.*?)\n```', code, re.DOTALL)
+        if not blocks:
+            # Try to write as single file
+            path = os.path.join(WORK_DIR, f"{base_name}.py")
+            with open(path, 'w') as f:
+                f.write(code)
+            written.append(path)
+        else:
+            for i, block in enumerate(blocks):
+                fname = f"{base_name}_{i}.py" if i > 0 else f"{base_name}.py"
+                path = os.path.join(WORK_DIR, fname)
+                with open(path, 'w') as f:
+                    f.write(block.strip() + "\n")
+                written.append(path)
+        return written
+    
+    def run_veritas_on_project(self) -> VeritasResult:
+        """Run Veritas on all .py files in WORK_DIR."""
+        all_errors = []
+        total_score = 0
+        file_count = 0
+        
+        for py_file in Path(WORK_DIR).glob("*.py"):
+            with open(py_file) as f:
+                code = f.read()
+            vr = self.veritas.run(code, str(py_file))
+            total_score += vr.score
+            file_count += 1
+            if not vr.passed:
+                all_errors.extend(vr.errors)
+        
+        if file_count == 0:
+            return VeritasResult(False, 0, {}, ["No Python files found"])
+        
+        avg_score = total_score / file_count
+        passed = avg_score >= 5 and len(all_errors) == 0
+        
+        return VeritasResult(
+            passed, 
+            int(avg_score),
+            {"avg_score": int(avg_score), "files_checked": file_count},
+            all_errors
+        )
+    
+    def run(self, task: str, verbose: bool = True) -> Dict[str, Any]:
+        if verbose:
+            print("=" * 55)
+            print(f"  SECOND BRAIN: {task[:50]}...")
+            print("=" * 55)
+        
+        self.load_model()
+        self.history = []
+        self.history.append(Step(thought=self.build_prompt(task)))
+        
+        for iteration in range(MAX_ITERATIONS):
+            if verbose:
+                print(f"\n  --- Iteration {iteration + 1} ---")
+            
+            # Generate
+            generated = self.generate_next()
+            tool_call = self.parse_tool_call(generated)
+            thought = self.extract_thought(generated, tool_call)
+            
+            self.add_to_history(Step(thought=thought, tool_call=tool_call, iteration=iteration))
+            
+            if verbose:
+                print(f"  THOUGHT: {thought[:100]}")
+                if tool_call:
+                    print(f"  TOOL: {json.dumps(tool_call)[:120]}")
+            
+            if not tool_call:
+                if verbose:
+                    print("  No tool call — retrying...")
+                self.add_to_history(Step(thought="I must emit a valid JSON tool call.", iteration=iteration))
+                continue
+            
+            if tool_call.get("tool") == "done":
+                if verbose:
+                    print("  TASK COMPLETE!")
+                # Final Veritas check
+                vr = self.run_veritas_on_project()
+                return {
+                    "success": vr.passed,
+                    "iterations": iteration + 1,
+                    "veritas_score": vr.score,
+                    "veritas_details": vr.details,
+                    "errors": vr.errors,
+                    "files": [str(p) for p in Path(WORK_DIR).glob("*.py")],
+                    "history": [s.__dict__ for s in self.history]
+                }
+            
+            # Execute tool
+            result = self.executor.execute(tool_call)
+            
+            # If writing code, also write to disk and run Veritas
+            if tool_call.get("tool") == "fs" and tool_call.get("action") == "write":
+                content = tool_call.get("content", "")
+                if content.strip():
+                    self.write_code_files(content, "output")
+                    vr = self.run_veritas_on_project()
+                    if not vr.passed:
+                        result = ToolResult(False, "", f"Veritas failed: {vr.errors}")
+            
+            self.history[-1].result = result
+            
+            if verbose:
+                if result.success:
+                    print(f"  RESULT: {result.result[:120]}")
+                else:
+                    print(f"  ERROR: {result.error[:120]}")
+        
+        if verbose:
+            print("  MAX ITERATIONS REACHED")
+        
+        vr = self.run_veritas_on_project()
+        return {
+            "success": False,
+            "iterations": MAX_ITERATIONS,
+            "veritas_score": vr.score,
+            "errors": ["Max iterations"] + vr.errors,
+            "files": [str(p) for p in Path(WORK_DIR).glob("*.py")],
+            "history": [s.__dict__ for s in self.history]
+        }
+
+
+# ═══════════════════════════════════════════════════════════
+# MOCK MODEL FOR TESTING (no weights needed)
+# ═══════════════════════════════════════════════════════════
+
+class MockModel:
+    """Deterministic mock that emits valid tool-call sequences for testing."""
+    
+    def __init__(self, tokenizer):
+        self.tok = tokenizer
+        self.call_count = 0
+        self.max_seq = 512
+    
+    def generate(self, ids, max_new=150, **kwargs):
+        self.call_count += 1
+        # Always return step-by-step sequence for first 3 calls to test the loop
+        if self.call_count == 1:
+            return self.tok.encode(
+                "THOUGHT: I will write a simple add function and verify it compiles.\n"
+                '{"tool": "fs", "action": "write", "path": "add.py", "content": "def add(a, b):\\n    return a + b\\n"}'
+            )
+        elif self.call_count == 2:
+            return self.tok.encode(
+                "THOUGHT: Now I will compile-check it.\n"
+                '{"tool": "compile", "path": "add.py"}'
+            )
+        else:
+            return self.tok.encode(
+                "THOUGHT: All checks passed.\n"
+                '{"tool": "done"}'
+            )
+
+
+# ═══════════════════════════════════════════════════════════
+# MAIN / TEST
+# ═══════════════════════════════════════════════════════════
+
+def test_with_mock():
+    """Test the loop with a mock model (no training needed)."""
+    print("=" * 55)
+    print("  SECOND BRAIN TEST (Mock Model)")
+    print("=" * 55)
+    
+    tok = BPETokenizer(vocab_size=4096)
+    tok.load(TOKENIZER_PATH)
+    
+    brain = SecondBrain(model=MockModel(tok), tokenizer=tok)
+    
+    test_tasks = [
+        "Create a Python function that adds two numbers and verify it compiles",
+    ]
+    
+    for task in test_tasks:
+        result = brain.run(task)
+        print(f"\n  SUCCESS: {result['success']}")
+        print(f"  ITERATIONS: {result['iterations']}")
+        print(f"  VERITAS: {result['veritas_score']}/6")
+        print(f"  FILES: {result['files']}")
+        if result['errors']:
+            print(f"  ERRORS: {result['errors']}")
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("task", nargs="*", help="Task to run")
+    parser.add_argument("--mock", action="store_true", help="Use mock model for testing")
+    args = parser.parse_args()
+    
+    if args.mock or not args.task:
+        test_with_mock()
+    else:
+        task = " ".join(args.task)
+        brain = SecondBrain()
+        result = brain.run(task)
+        print(f"\n  FINAL: {'SUCCESS' if result['success'] else 'FAILED'}")
+        print(f"  Veritas: {result['veritas_score']}/6")
+        if result['errors']:
+            print(f"  Errors: {result['errors']}")