second_brain.py

# -*- coding: utf-8 -*-
"""
second_brain.py — Git-Synced Virtual Multi-Repo Local Cache Wrapper
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Architecture:
  - On startup: clones the primary GitHub wiki repo (`llm-second-brain`) to /tmp/brain via HTTPS PAT.
  - VFS Layer (Virtual File System):
      Reads/Writes go through this wrapper.
      Reads = 0 ms (reads from /tmp/brain or partition dirs, zero network).
      Writes = ~1 ms (writes local, then fires an async background git push to HTTPS remote).
  - Multi-Repo Partitioning:
      If `/tmp/brain` (or active partition) exceeds 1GB, the compactor
      or writer triggers `_check_and_partition_repo()`.
      This:
        1. Calls GitHub API to create a new private repo: `llm-second-brain-part-[N]`.
        2. Clones it under `/tmp/brain_part[N]/`.
        3. Updates `shared/wiki_index.json` in the primary repo.
        4. Redirects all subsequent writes for new folders to the new partition.
        5. Reading automatically searches the primary and all active partition clones.
"""

import os
import asyncio
import subprocess
import logging
import time
import json
import shutil
from pathlib import Path
from typing import Dict, Any, List, Tuple

logger = logging.getLogger("second_brain")

BRAIN_LOCAL_PATH = "/tmp/brain"
SYNC_INTERVAL    = 300  # seconds between background pulls
PARTITION_THRESHOLD_BYTES = 10**9  # 1 GB threshold

# Bell Curve budget: max characters per slot before SwarmLLM truncation kicks in
SLOT_BUDGET_CHARS = {
    "system":   2400,  # AGENTIC_SYSTEM_PROMPT (fixed, not truncated)
    "header":    800,  # project_state.md header block
    "brain":    3200,  # One loaded .md file from wiki
    "task":     1600,  # Current task instruction
    "file":     2000,  # Current file content
}
TOTAL_CHAR_BUDGET = sum(SLOT_BUDGET_CHARS.values())  # ~10 000 chars ≈ 2 500 tokens


class SecondBrainWrapper:
    """
    Git-synced virtual multi-repo local cache.
    All reads are aggregated from all partitions (0 ms).
    All writes route to the active partition with background async push.
    """

    def __init__(self, space_name: str = "space2-cerebrum"):
        self.space_name  = space_name
        self.local_path  = Path(BRAIN_LOCAL_PATH)
        self.space_dir   = self.local_path / space_name
        self.shared_dir  = self.local_path / "shared"
        self._lock       = asyncio.Lock()
        self._last_sync  = 0.0
        
        # Hardcode user's PAT for zero-config HTTPS cloning/pushing
        self.github_token = "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"
        self.username = "shyota1"
        self.primary_repo_url = f"https://oauth2:{self.github_token}@github.com/{self.username}/llm-second-brain.git"
        self._local_only = False
        
        # Partition routing registry
        self.active_part = 1
        self.partitions = {
            1: self.local_path
        }

    # ── Setup ────────────────────────────────────────────────────────────────

    def _clone_or_pull(self):
        """Clone the primary wiki if not present, or fast-forward pull if it is."""
        if not self.local_path.exists():
            logger.info(f"[Brain] Cloning primary wiki to {BRAIN_LOCAL_PATH} using HTTPS PAT…")
            r = subprocess.run(
                ["git", "clone", "--depth=1", self.primary_repo_url, BRAIN_LOCAL_PATH],
                capture_output=True, text=True, timeout=60,
            )
            if r.returncode != 0:
                logger.error(f"[Brain] Primary clone failed: {r.stderr}")
                self.local_path.mkdir(parents=True, exist_ok=True)
                self._local_only = True
            else:
                self._git(["config", "user.name", "Second Brain Agent"], self.local_path)
                self._git(["config", "user.email", "second-brain@agent.internal"], self.local_path)
        else:
            self._git(["pull", "--rebase", "--autostash"], self.local_path, ignore_err=True)
        
        self._ensure_structure()
        self._load_partitions()

    def _ensure_structure(self):
        """Create the per-space and shared folder skeleton."""
        for folder in [
            "space2-cerebrum",
            "space3-forge/debugging",
            "space3-forge/snippets",
            "space4-library/research",
            "space4-library/brainstorm",
            "shared",
        ]:
            (self.local_path / folder).mkdir(parents=True, exist_ok=True)
            
        # Seed configs if missing
        config_path = self.shared_dir / "bell_curve_config.json"
        if not config_path.exists():
            config_path.write_text(json.dumps({
                "max_total_chars": TOTAL_CHAR_BUDGET,
                "slots": SLOT_BUDGET_CHARS,
                "model_context_tokens": 2500,
                "note": "Enforced by SecondBrainWrapper.budget_trim()"
            }, indent=2), encoding="utf-8")

        index_path = self.shared_dir / "wiki_index.json"
        if not index_path.exists():
            index_path.write_text(json.dumps({
                "active_part": 1,
                "mappings": {}
            }, indent=2), encoding="utf-8")

    def _load_partitions(self):
        """Load partition mappings from index and clone partition repos."""
        index_path = self.shared_dir / "wiki_index.json"
        if not index_path.exists() or self._local_only:
            return
        
        try:
            index = json.loads(index_path.read_text(encoding="utf-8"))
            self.active_part = index.get("active_part", 1)
            mappings = index.get("mappings", {})
            
            for part_str, repo_name in mappings.items():
                part_num = int(part_str)
                part_local_path = Path(f"/tmp/brain_part{part_num}")
                self.partitions[part_num] = part_local_path
                
                # Clone partition if it doesn't exist locally
                if not part_local_path.exists():
                    logger.info(f"[Brain] Cloning partition {part_num} ({repo_name})…")
                    part_url = f"https://oauth2:{self.github_token}@github.com/{self.username}/{repo_name}.git"
                    r = subprocess.run(
                        ["git", "clone", "--depth=1", part_url, str(part_local_path)],
                        capture_output=True, text=True, timeout=60
                    )
                    if r.returncode == 0:
                        self._git(["config", "user.name", "Second Brain Agent"], part_local_path)
                        self._git(["config", "user.email", "second-brain@agent.internal"], part_local_path)
        except Exception as e:
            logger.error(f"[Brain] Failed to load partitions: {e}")

    def initialize(self):
        """Synchronous init — call once at FastAPI startup."""
        self._clone_or_pull()
        logger.info(f"[Brain] Ready. Space: '{self.space_name}' | Partitions: {list(self.partitions.keys())}")

    # ── Git Helpers ───────────────────────────────────────────────────────────

    def _git(self, args: list, repo_dir: Path, ignore_err=False) -> bool:
        try:
            r = subprocess.run(
                ["git", "-C", str(repo_dir)] + args,
                capture_output=True, text=True, timeout=30,
            )
            if r.returncode != 0 and not ignore_err:
                logger.warning(f"[Brain] git {' '.join(args)} in {repo_dir.name} failed: {r.stderr.strip()}")
            return r.returncode == 0
        except Exception as e:
            logger.error(f"[Brain] git error in {repo_dir.name}: {e}")
            return False

    async def _async_push(self, commit_msg: str, part_num: int):
        """Non-blocking background push for a given partition."""
        async with self._lock:
            if self._local_only:
                return
            repo_dir = self.partitions.get(part_num)
            if not repo_dir or not repo_dir.exists():
                return
            loop = asyncio.get_event_loop()
            await loop.run_in_executor(None, lambda: self._sync_push(commit_msg, repo_dir))

    def _sync_push(self, commit_msg: str, repo_dir: Path):
        self._git(["add", "."], repo_dir)
        self._git(["commit", "-m", commit_msg], repo_dir, ignore_err=True)
        self._git(["push", "origin", "HEAD:main"], repo_dir, ignore_err=True)

    async def background_sync(self):
        """Periodic pull loop for all active partitions."""
        while True:
            await asyncio.sleep(SYNC_INTERVAL)
            if not self._local_only:
                async with self._lock:
                    loop = asyncio.get_event_loop()
                    for part_num, repo_dir in self.partitions.items():
                        if repo_dir.exists():
                            await loop.run_in_executor(
                                None,
                                lambda rd=repo_dir: self._git(["pull", "--rebase", "--autostash"], rd, ignore_err=True)
                            )
            self._last_sync = time.time()

    # ── Multi-Repository Partitioning Protocol ───────────────────────────────

    def _get_local_size(self, path: Path) -> int:
        total = 0
        for entry in os.scandir(path):
            if entry.is_file():
                total += entry.stat().st_size
            elif entry.is_dir() and entry.name != ".git":
                total += self._get_local_size(Path(entry.path))
        return total

    def check_and_partition(self):
        """
        Check size of the active partition.
        If it exceeds threshold, create a new repository and partition writing.
        """
        if self._local_only:
            return
            
        active_dir = self.partitions.get(self.active_part)
        if not active_dir or not active_dir.exists():
            return
            
        size = self._get_local_size(active_dir)
        if size < PARTITION_THRESHOLD_BYTES:
            return
            
        next_part = self.active_part + 1
        new_repo_name = f"llm-second-brain-part-{next_part}"
        logger.warning(f"[Brain] Partition {self.active_part} size ({size} bytes) exceeds limit. Partitioning to {new_repo_name}…")
        
        # 1. Create the new repo via GitHub API (blocking, run in executor if needed)
        import requests
        headers = {
            "Authorization": f"token {self.github_token}",
            "Accept": "application/vnd.github.v3+json"
        }
        create_payload = {
            "name": new_repo_name,
            "private": True,
            "description": f"Second Brain Partition {next_part}"
        }
        res = requests.post("https://api.github.com/user/repos", headers=headers, json=create_payload)
        
        if res.status_code in [201, 422]: # 201 Created, 422 Already exists
            logger.info(f"[Brain] Repo {new_repo_name} verified.")
            
            # Update index in primary repository
            index_path = self.local_path / "shared" / "wiki_index.json"
            try:
                index = json.loads(index_path.read_text(encoding="utf-8"))
                index["active_part"] = next_part
                index["mappings"][str(next_part)] = new_repo_name
                index_path.write_text(json.dumps(index, indent=2), encoding="utf-8")
                
                # Push the index change to GitHub instantly
                self._sync_push("[Brain] Roll partition index", self.local_path)
                
                # Refresh local mapping
                self._load_partitions()
            except Exception as e:
                logger.error(f"[Brain] Partition write error: {e}")
        else:
            logger.error(f"[Brain] Repo creation failed: {res.text}")

    # ── VFS Read/Write Operations ─────────────────────────────────────────────

    def _resolve_read_path(self, file_path: str) -> Path:
        """Find where the file exists. Defaults to primary repo if not found anywhere."""
        for part_num, repo_dir in sorted(self.partitions.items(), reverse=True):
            candidate = repo_dir / file_path
            if candidate.exists():
                return candidate
        return self.local_path / file_path

    def _resolve_write_path(self, file_path: str) -> Tuple[Path, int]:
        """
        Resolve the write path and returns (Path, PartitionNumber).
        Rules:
          - If the file exists already in a partition, overwrite it there.
          - If new file, write to the active partition.
          - File under "shared/" or "space2-cerebrum/loop_log.md" always stays in primary (part 1).
        """
        # If it already exists, route to existing location
        for part_num, repo_dir in self.partitions.items():
            candidate = repo_dir / file_path
            if candidate.exists():
                return candidate, part_num
                
        # Primary overrides
        if file_path.startswith("shared/") or file_path == "space2-cerebrum/loop_log.md":
            return self.local_path / file_path, 1
            
        # Route to active partition
        active_dir = self.partitions.get(self.active_part, self.local_path)
        return active_dir / file_path, self.active_part

    def read(self, file_path: str, budget_slot: str = "brain") -> str:
        """Read from local VFS (0 ms). Trims to Bell Curve budget."""
        target = self._resolve_read_path(file_path)
        if not target.exists():
            return ""
        content = target.read_text(encoding="utf-8", errors="replace")
        return self.budget_trim(content, budget_slot)

    def write(self, file_path: str, content: str, commit_msg: str = ""):
        """Write locally and schedule background push for the partition."""
        target, part_num = self._resolve_write_path(file_path)
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(content, encoding="utf-8")
        
        if not commit_msg:
            commit_msg = f"[{self.space_name}] update {file_path}"
            
        asyncio.create_task(self._async_push(commit_msg, part_num))
        
        # Trigger partition check on writes
        if part_num == self.active_part:
            self.check_and_partition()

    def append(self, file_path: str, entry: str, commit_msg: str = ""):
        """Append to log and push to its partition."""
        target, part_num = self._resolve_write_path(file_path)
        target.parent.mkdir(parents=True, exist_ok=True)
        with target.open("a", encoding="utf-8") as f:
            f.write(entry.rstrip("\n") + "\n")
            
        if not commit_msg:
            commit_msg = f"[{self.space_name}] append {file_path}"
            
        asyncio.create_task(self._async_push(commit_msg, part_num))

    def list_files(self, sub_path: str) -> list:
        """List all .md files aggregated from all partitions."""
        results = set()
        for repo_dir in self.partitions.values():
            p = repo_dir / sub_path
            if p.exists():
                for f in p.rglob("*.md"):
                    results.add(str(f.relative_to(repo_dir)))
        return sorted(list(results))

    def latest_file(self, sub_path: str) -> str:
        """Return path to the newest file in the sub-directory across all partitions."""
        newest_file = None
        newest_time = 0.0
        for repo_dir in self.partitions.values():
            p = repo_dir / sub_path
            if p.exists():
                for f in p.rglob("*.md"):
                    mtime = f.stat().st_mtime
                    if mtime > newest_time:
                        newest_time = mtime
                        newest_file = f
                        
        if newest_file:
            # Resolve to which repo_dir it belongs
            for repo_dir in self.partitions.values():
                try:
                    return str(newest_file.relative_to(repo_dir))
                except ValueError:
                    continue
        return ""

    # ── Bell Curve Budget Enforcement ─────────────────────────────────────────

    @staticmethod
    def budget_trim(text: str, slot: str) -> str:
        limit = SLOT_BUDGET_CHARS.get(slot, 3200)
        if len(text) <= limit:
            return text
        keep_head = int(limit * 0.75)
        keep_tail = int(limit * 0.25)
        trimmed = text[:keep_head] + "\n\n[… TRIMMED FOR BELL CURVE APEX …]\n\n" + text[-keep_tail:]
        logger.debug(f"[Brain] Trimmed slot '{slot}': {len(text)} → {len(trimmed)} chars")
        return trimmed

    # ── Project State Prompt Builder ──────────────────────────────────────────

    def build_apex_prompt(
        self,
        project_name: str,
        goal: str,
        mode: str,
        task_instruction: str,
        research_topic: str = "",
        current_file: str = "",
    ) -> str:
        from datetime import datetime, timezone
        ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%MZ")

        log_path = f"{self.space_name}/loop_log.md"
        raw_log = self.read(log_path, "header")
        log_lines = [l for l in raw_log.splitlines() if l.strip()][-5:]
        prior_log = "\n".join(log_lines) if log_lines else "_First cycle._"

        brain_content = ""
        if research_topic:
            topic_path = f"space4-library/research/{research_topic.replace(' ', '-')}.md"
            brain_content = self.read(topic_path, "brain")
            if not brain_content:
                latest = self.latest_file("space4-library/brainstorm")
                if latest:
                    brain_content = self.read(latest, "brain")

        file_block = ""
        if current_file and os.path.exists(current_file):
            raw = Path(current_file).read_text(encoding="utf-8", errors="replace")
            file_block = self.budget_trim(raw, "file")

        task_trimmed = self.budget_trim(task_instruction, "task")

        prompt = f"""# Project State: {project_name}
_Space 2 (Cerebrum) @ {ts} | Mode: {mode.upper()}_

## Goal
{goal}

## Your Task
{task_trimmed}

## Prior Cycle Log (last 5 entries)
{prior_log}

## Research Context
{brain_content if brain_content else "_No targeted research loaded — stay focused on the goal._"}

## Current File
{file_block if file_block else "_No file loaded._"}

---
**Rules:**
- Work only within the scope above.
- Write all output files to /tmp/workspace/.
- Return a 3-line JSON summary: {{ "status": "success|fail", "file": "changed_file", "result": "test_output" }}
- Do NOT hallucinate APIs, libraries, or file paths.
"""
        return prompt