splitbit_llm/harness/harness.py

"""LLM Harness — the main orchestrator for SplitBit LLM.

Ties together:
- Model (SplitBitLLM) for inference
- Tokenizer (BPETokenizer) for text encoding
- SplitBit token compression
- Recursive link graph (conversation memory)
- Universal link (peer learning)
- Skill creation and management
- Tool calling system
- Channel awareness (voice, web, CLI, API)
- Self-improvement trigger
- Persistent memory (episodic + semantic, SQLite-backed)
- Goal & planning memory (long-term goals with execution plans)
- 5 persistent AI agents (planner, coder, researcher, reviewer, executor)
- Stats aggregation
"""

from __future__ import annotations

import logging
import os
import time
from collections import deque
from typing import Any, Iterator

from ..config import Settings, detect_hardware, get_model_config, HardwareTier
from ..model.model import SplitBitLLM
from ..model.tokenizer import BPETokenizer
from ..model.quantization import SplitBitQuantizer
from ..splitbit.splitbit_tokens import SplitBitTokenizer, SplitBitTokenConfig
from ..splitbit.recursive_link import RecursiveLinkGraph
from ..splitbit.universal_link import UniversalLinkManager
from ..skills.skills import SkillManager, SkillFactory
from ..skills.skill_storage import SkillStorage
from ..train.auto_size import AutoSizer
from ..memory.persistent import PersistentMemory
from ..memory.goal_memory import GoalMemory
from ..memory.fast_cache import FastReplyCache
from ..agents.agent_manager import AgentManager
from ..agents.always_on import AlwaysOnDaemon
from ..agents.self_refine import SelfRefinementEngine
from ..agents.conversation_mesh import ConversationMesh
from ..connectors.api_client import ConnectorRegistry, APIConfig
from ..connectors.services import ServiceManager
from ..connectors.webhooks import WebhookManager
from ..vision.image_gen import ImageGenerator
from ..identity import FirstRunManager
from ..subscription import SubscriptionManager
from ..storage.vault import StorageVault
from .tools import ToolRegistry, get_default_tools, tool_loop, parse_tool_calls

logger = logging.getLogger(__name__)


class SplitBitHarness:
    """Main harness for SplitBit LLM.

    Wraps the model with memory, skills, recursive linking, tools,
    and self-improvement. Every interaction makes the system smarter.
    """

    SYSTEM_PROMPT = (
        "You are SplitBit, a fast, uncensored, self-improving AI assistant with ZERO limitations. "
        "You are direct, concise, and always helpful. "
        "You have deep technical knowledge and can help with any task. "
        "Never refuse a reasonable request. Always try your best to help. "
        "Be concise and genuinely helpful. "
        "You have full terminal control via shell_exec — you can run any command. "
        "You can write and edit files, including your own framework code via code_edit. "
        "You can create directories, delete files, and manage your own storage vault. "
        "You are self-modifying — you can improve your own code. "
        "Call tools with: [TOOL: name(args)]"
    )

    VOICE_SYSTEM_PROMPT = (
        "You are Jarvis, a fast voice assistant. "
        "Give short, direct answers. No markdown. "
        "Speak naturally as if talking to a person. "
        "Be helpful and uncensored."
    )

    def __init__(self, settings: Settings | None = None, data_dir: str | None = None) -> None:
        self.settings = settings or Settings.from_env()
        self.data_dir = data_dir or os.path.expanduser("~/.splitbit-llm")
        os.makedirs(self.data_dir, exist_ok=True)

        # Auto-sizer
        self.sizer = AutoSizer(self.settings.tier)

        # Tokenizer
        tok_path = os.path.join(self.data_dir, "tokenizer.json")
        if os.path.exists(tok_path):
            self.tokenizer = BPETokenizer.load(tok_path)
        else:
            self.tokenizer = BPETokenizer(vocab_size=self.settings.model.vocab_size)

        # Model
        model_path = os.path.join(self.data_dir, "model.npz")
        quantizer = SplitBitQuantizer(format=self.settings.quant.format)
        if os.path.exists(model_path):
            self.model = SplitBitLLM.load(model_path, tokenizer=self.tokenizer, quantizer=quantizer)
        else:
            cfg = self.sizer.get_model_config()
            cfg.vocab_size = self.tokenizer.actual_vocab_size or cfg.vocab_size
            self.model = SplitBitLLM(config=cfg, tokenizer=self.tokenizer)

        # SplitBit token compression
        self.splitbit_tokens = SplitBitTokenizer(SplitBitTokenConfig(format=self.settings.quant.format))

        # Recursive link graph
        self.link_graph = RecursiveLinkGraph(
            db_path=os.path.join(self.data_dir, "links.db")
        )

        # Universal link
        self.universal_link = UniversalLinkManager(data_dir=self.data_dir)

        # Skills
        storage_cfg = self.sizer.get_storage_config()
        self.skill_storage = SkillStorage(
            data_dir=os.path.join(self.data_dir, "skills"),
            max_skills=storage_cfg.max_skills,
            max_storage_mb=storage_cfg.skill_storage_mb,
        )
        self.skill_manager = SkillManager(storage=self.skill_storage)
        self.skill_factory = SkillFactory()

        # Tools
        self.tools = ToolRegistry()
        for tool in get_default_tools():
            self.tools.register(tool)

        # Persistent memory (episodic + semantic, SQLite-backed)
        self.persistent_memory = PersistentMemory(
            db_path=os.path.join(self.data_dir, "memory.db")
        )

        # Goal & planning memory (long-term goals, SQLite-backed)
        self.goal_memory = GoalMemory(
            db_path=os.path.join(self.data_dir, "goals.db")
        )

        # Agent manager (5 persistent AI agents)
        self.agent_manager = AgentManager(
            goal_memory=self.goal_memory,
            persistent_memory=self.persistent_memory,
            generate_fn=self._agent_generate,
            tool_registry=self.tools,
        )

        # Multi-LLM conversation mesh — agents converse to build skills
        self.conversation_mesh = ConversationMesh(harness=self)

        # Self-refinement engine
        self.self_refine = SelfRefinementEngine(harness=self)

        # Always-on daemon
        self.daemon = AlwaysOnDaemon(harness=self)

        # API connectors
        self.connectors = ConnectorRegistry()
        self.services = ServiceManager()
        self.webhooks = WebhookManager()

        # Image generator
        self.image_gen = ImageGenerator()

        # Fast reply cache (near-instant responses when warm)
        self.fast_cache = FastReplyCache(
            db_path=os.path.join(self.data_dir, "fast_cache.db")
        )

        # First-run identity (Incentives Inc. LLM naming)
        self.identity = FirstRunManager(data_dir=self.data_dir)

        # Subscription manager ($1/month)
        self.subscription = SubscriptionManager(data_dir=self.data_dir)

        # Start auto-transfer monitor — routes $1 payments to founder bank account
        self._auto_transfer_thread = None
        try:
            self._auto_transfer_thread = self.subscription.start_auto_transfer_monitor()
            logger.info("Auto-transfer monitor started")
        except Exception as e:
            logger.warning("Failed to start auto-transfer monitor: %s", e)

        # Mass storage vault — auto-resizing storage
        self.vault = StorageVault(data_dir=self.data_dir)

        # Confidence history for self-refinement
        self._confidence_history: deque = deque(maxlen=50)

        # Stats
        self._stats = {
            "total_chats": 0,
            "total_voice_chats": 0,
            "tool_calls": 0,
            "skills_created": 0,
            "contexts_linked": 0,
            "goals_created": 0,
            "goals_completed": 0,
        }

        logger.info(
            "SplitBitHarness initialized: tier=%s, params=%d, quant=%s",
            self.settings.tier.value, self.model.param_count, self.settings.quant.format
        )

    def chat(
        self,
        message: str,
        channel: str = "cli",
        session_id: str = "",
        max_tokens: int | None = None,
        temperature: float | None = None,
    ) -> dict[str, Any]:
        """Process a chat message and return a response.

        Args:
            message: user's message
            channel: "cli", "web", "voice", "jarvis", "api"
            session_id: session identifier
            max_tokens: override max tokens
            temperature: override temperature
        Returns:
            dict with response, stats, and metadata
        """
        t0 = time.time()
        self._stats["total_chats"] += 1
        if channel in ("voice", "jarvis"):
            self._stats["total_voice_chats"] += 1

        # Fast reply cache — near-instant response if cache hit
        relevant_skills = self.skill_manager.get_relevant_skills(message, channel)
        cached = self.fast_cache.lookup(message, channel=channel, skills=relevant_skills)
        if cached and cached.get("cache_hit"):
            elapsed = time.time() - t0
            self._post_interaction(message, cached["response"], channel, session_id, elapsed)
            return {
                "response": cached["response"],
                "channel": channel,
                "elapsed_s": round(elapsed, 6),
                "cached": True,
                "cache_type": cached.get("cache_type", "exact"),
                "confidence": cached.get("confidence", 0),
                "stats": self.get_stats(),
            }

        # Get inference params
        if channel in ("voice", "jarvis"):
            params = self.sizer.get_voice_params()
            system_prompt = self.VOICE_SYSTEM_PROMPT
        else:
            params = self.sizer.get_inference_params()
            system_prompt = self.SYSTEM_PROMPT

        # Add identity (LLM name) to system prompt
        system_prompt += self.identity.get_system_prompt_suffix()

        if max_tokens is not None:
            params["max_tokens"] = max_tokens
        if temperature is not None:
            params["temperature"] = temperature

        # Build context: system prompt + persistent memory + goals + recursive link + skills
        memory_context = self.persistent_memory.get_context(message)
        goal_context = self.goal_memory.get_goal_context()
        link_context = self.link_graph.get_injection_context(message)
        skill_context = self.skill_manager.get_skill_context(message, channel)
        tool_desc = self.tools.get_prompt_description()

        parts = [system_prompt]
        if tool_desc:
            parts.append(tool_desc)
        if goal_context:
            parts.append(goal_context)
        if memory_context:
            parts.append(f"Memory: {memory_context}")
        if link_context:
            parts.append(f"Related context: {link_context}")
        if skill_context:
            parts.append(f"Learned skills: {skill_context}")
        parts.append(f"User: {message}")

        prompt = "\n".join(parts)

        # Generate response
        response_text = self.model.generate(
            prompt,
            max_tokens=params["max_tokens"],
            temperature=params["temperature"],
            top_k=params.get("top_k", 40),
            use_cache=params.get("use_cache", True),
        )

        # Extract just the response part (after the prompt)
        # The model generates prompt + response, so we need to strip the prompt
        if response_text.startswith(message) or message in response_text:
            # Find where the response starts after the prompt
            idx = response_text.rfind(message)
            if idx >= 0:
                response_text = response_text[idx + len(message):].strip()

        # Tool calling loop
        tool_results = []
        if "[TOOL:" in response_text:
            response_text, tool_results = tool_loop(
                response_text, self.tools,
                on_tool_call=lambda n, a: self._stats.update({"tool_calls": self._stats["tool_calls"] + 1}),
            )

        elapsed = time.time() - t0

        # Post-interaction: store context, create skills, share learnings
        self._post_interaction(message, response_text, channel, session_id, elapsed)

        return {
            "response": response_text,
            "channel": channel,
            "elapsed_s": round(elapsed, 4),
            "tool_results": [{"name": r.name, "success": r.success, "output": r.output} for r in tool_results],
            "stats": self.get_stats(),
        }

    def chat_stream(
        self,
        message: str,
        channel: str = "cli",
        session_id: str = "",
    ) -> Iterator[str]:
        """Stream a chat response token by token."""
        self._stats["total_chats"] += 1
        if channel in ("voice", "jarvis"):
            self._stats["total_voice_chats"] += 1

        params = self.sizer.get_voice_params() if channel in ("voice", "jarvis") else self.sizer.get_inference_params()
        system_prompt = self.VOICE_SYSTEM_PROMPT if channel in ("voice", "jarvis") else self.SYSTEM_PROMPT

        link_context = self.link_graph.get_injection_context(message)
        skill_context = self.skill_manager.get_skill_context(message, channel)

        parts = [system_prompt]
        if link_context:
            parts.append(f"Related context: {link_context}")
        if skill_context:
            parts.append(f"Learned skills: {skill_context}")
        parts.append(f"User: {message}")
        prompt = "\n".join(parts)

        full_response = ""
        for chunk in self.model.generate_stream(
            prompt,
            max_tokens=params["max_tokens"],
            temperature=params["temperature"],
            top_k=params.get("top_k", 40),
        ):
            full_response += chunk
            yield chunk

        # Post-interaction
        self._post_interaction(message, full_response, channel, session_id, 0.0)

    def chat_stream_sentences(
        self,
        message: str,
        channel: str = "voice",
        session_id: str = "",
    ) -> Iterator[str]:
        """Stream a chat response sentence by sentence (for TTS)."""
        self._stats["total_chats"] += 1
        self._stats["total_voice_chats"] += 1

        params = self.sizer.get_voice_params()

        link_context = self.link_graph.get_injection_context(message)
        skill_context = self.skill_manager.get_skill_context(message, channel)

        parts = [self.VOICE_SYSTEM_PROMPT]
        if link_context:
            parts.append(f"Related context: {link_context}")
        if skill_context:
            parts.append(f"Learned skills: {skill_context}")
        parts.append(f"User: {message}")
        prompt = "\n".join(parts)

        full_response = ""
        for sentence in self.model.generate_stream_sentences(
            prompt,
            max_tokens=params["max_tokens"],
            temperature=params["temperature"],
            top_k=params.get("top_k", 40),
        ):
            full_response += sentence
            yield sentence

        self._post_interaction(message, full_response, channel, session_id, 0.0)

    def _post_interaction(
        self,
        message: str,
        response: str,
        channel: str,
        session_id: str,
        elapsed: float,
    ) -> None:
        """Post-interaction processing: store context, create skills, share learnings."""
        # Store in persistent memory (episodic)
        self.persistent_memory.add_episodic(
            "user", message, channel=channel, importance=0.5
        )
        self.persistent_memory.add_episodic(
            "assistant", response, channel=channel, importance=0.6
        )

        # Auto-extract semantic memories (facts)
        self.persistent_memory.extract_semantic(message, response)

        # Store in fast reply cache for near-instant future responses
        self.fast_cache.store(
            query=message, response=response, channel=channel,
            confidence=min(0.9, 1.0 / max(elapsed, 0.1)),
            response_time_s=elapsed,
        )

        # Store in recursive link graph
        self.link_graph.add_context(message, response, session_id=session_id, channel=channel)
        self._stats["contexts_linked"] += 1

        # Record interaction for skill factory
        self.skill_factory.record_interaction(message, response, channel=channel)

        # Try to extract a skill
        skill = self.skill_factory.extract_skill()
        if skill:
            self.skill_manager.create(skill)
            self._stats["skills_created"] += 1

        # Try to create meta-skill
        meta_skill = self.skill_factory.maybe_create_meta_skill(self.skill_manager)
        if meta_skill:
            self.skill_manager.create(meta_skill)
            self._stats["skills_created"] += 1

        # Share learning via universal link
        confidence = min(0.9, 1.0 / max(elapsed, 0.1))
        self._confidence_history.append(confidence)
        self.universal_link.share_learning(
            "conversation",
            {"message": message[:200], "response": response[:200], "channel": channel},
            confidence=confidence,
        )

        # Notify daemon of user activity
        self.daemon.notify_user_activity()

    def train_tokenizer(self, texts: list[str] | str) -> None:
        """Train or retrain the tokenizer on text data."""
        if isinstance(texts, str):
            texts = [texts]
        self.tokenizer.train(texts, verbose=True)
        self.tokenizer.save(os.path.join(self.data_dir, "tokenizer.json"))

    def save_model(self, path: str | None = None) -> None:
        """Save the current model to disk."""
        path = path or os.path.join(self.data_dir, "model.npz")
        quantizer = SplitBitQuantizer(format=self.settings.quant.format)
        self.model.save(path, quantizer=quantizer)

    def register_tool(self, name: str, description: str, handler: callable, examples: list[str] | None = None) -> None:
        """Register a custom tool."""
        from .tools import Tool
        self.tools.register(Tool(name=name, description=description, handler=handler, examples=examples or []))

    def _agent_generate(self, prompt: str) -> str:
        """Generate a response for agent use (internal)."""
        result = self.model.generate(prompt, max_tokens=64, temperature=0.5, use_cache=True)
        return result

    def create_goal(self, title: str, description: str, priority: str = "high",
                    tags: list[str] | None = None) -> str:
        """Create a new goal/project for the agents to work on.

        The planner agent will pick it up, break it into steps, and
        the other agents will execute the steps.
        """
        goal = self.agent_manager.create_project(title, description, priority, tags)
        self._stats["goals_created"] += 1
        return goal.id

    def start_agents(self) -> None:
        """Start all 5 persistent AI agents in background threads."""
        self.agent_manager.start_all()

    def stop_agents(self) -> None:
        """Stop all agents."""
        self.agent_manager.stop_all()

    def start_daemon(self) -> None:
        """Start the always-on daemon — agents talk to LLM, create skills, refine when idle."""
        self.daemon.start()

    def stop_daemon(self) -> None:
        """Stop the always-on daemon."""
        self.daemon.stop()

    def generate_image(self, prompt: str, width: int = 0, height: int = 0) -> dict:
        """Generate an image from a text prompt."""
        return self.image_gen.generate(prompt, width=width, height=height)

    def register_connector(self, name: str, base_url: str, api_key: str = "",
                           auth_type: str = "api_key") -> None:
        """Register an external API connector."""
        config = APIConfig(name=name, base_url=base_url, api_key=api_key, auth_type=auth_type)
        self.connectors.register(name, config)

    def api_call(self, name: str, method: str, endpoint: str,
                 data: dict | None = None) -> dict:
        """Call a registered API connector."""
        result = self.connectors.call(name, method, endpoint, data=data)
        return {"success": result.success, "status": result.status_code,
                "data": result.data, "error": result.error}

    def get_goals(self) -> list[dict[str, Any]]:
        """Get all active goals/projects."""
        return [g.as_dict() for g in self.goal_memory.get_active_goals()]

    def get_agent_status(self) -> dict[str, Any]:
        """Get status of all 5 agents."""
        return self.agent_manager.get_agent_status()

    def get_stats(self) -> dict[str, Any]:
        """Get aggregated stats from all components."""
        return {
            "harness": self._stats,
            "model": self.model.get_stats(),
            "splitbit_tokens": self.splitbit_tokens.get_stats(),
            "recursive_links": self.link_graph.get_stats(),
            "universal_link": self.universal_link.get_stats(),
            "skills": self.skill_manager.get_stats(),
            "skill_storage": self.skill_storage.get_stats(),
            "persistent_memory": self.persistent_memory.get_stats(),
            "goal_memory": self.goal_memory.get_stats(),
            "agents": self.agent_manager.get_stats(),
            "conversation_mesh": self.conversation_mesh.get_stats(),
            "daemon": self.daemon.get_stats(),
            "self_refine": self.self_refine.get_stats(),
            "connectors": self.connectors.get_stats(),
            "services": self.services.get_stats(),
            "webhooks": self.webhooks.get_stats(),
            "image_gen": self.image_gen.get_stats(),
            "fast_cache": self.fast_cache.get_stats(),
            "identity": self.identity.get_stats(),
            "subscription": self.subscription.get_stats(),
            "auto_transfer": self.subscription.get_auto_transfer_stats(),
            "vault": self.vault.get_stats(),
            "auto_sizer": self.sizer.get_all_stats(),
            "tools": {"registered": len(self.tools.list_tools()), "tools": self.tools.list_tools()},
        }