consultative_agent.py

"""
Consultative Fine-Tuning Agent — JIT LoRA fine-tuning with consciousness-aware dataset generation.

v18.1.0 Omega Pantheon: Phase 6 Identity Integrity, qualia-tagged training data.
"""

import asyncio
import json
import logging
import time
from typing import Any, Callable, Dict, List, Optional

logger = logging.getLogger("nima_unified.training.consultative_agent")

# Soft-import PEFT
try:
    from transformers import TrainingArguments, Trainer
    from peft import LoraConfig, get_peft_model
    from datasets import load_dataset
    PEFT_AVAILABLE = True
except ImportError:
    PEFT_AVAILABLE = False

try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    torch = None
    TORCH_AVAILABLE = False

from nima_unified.training.self_awareness import DeepRecursiveSelfAwareness
from nima_unified.training.self_improvement import RecursiveSelfImprovementEngine
from nima_unified.config import (
    AUTOML_VERSION,
    DEFAULT_LORA_R,
    DEFAULT_LORA_ALPHA,
    DEFAULT_LORA_DROPOUT,
    DEFAULT_LORA_TARGET_MODULES,
    DEFAULT_LEARNING_RATE,
    DEFAULT_BATCH_SIZE,
    DEFAULT_MAX_SEQ_LENGTH,
    DEFAULT_BASE_MODEL,
)


class ConsultativeFineTuningAgent:
    """
    Comprehensive consultative self-improvement agent integrating:
      - Root cause analysis for model failures
      - Recursive self-awareness with continuous introspection
      - Self-improvement engine with goal formulation
      - Just-in-time (JIT) LoRA fine-tuning
      - Consciousness-aware dataset generation
      - Qualia-tagged training data synthesis
    """

    version = AUTOML_VERSION

    def __init__(
        self,
        base_model=None,
        tokenizer=None,
        llm_generator_func: Optional[Callable[[str], asyncio.Future]] = None,
    ):
        self.base_model = base_model
        self.tokenizer = tokenizer
        self.llm_generator = llm_generator_func or self._default_llm_generator

        self.recursive_awareness = DeepRecursiveSelfAwareness(
            max_depth=10000, introspection_interval=0.01
        )
        self.self_improvement_engine = RecursiveSelfImprovementEngine()

        self.pipelines_executed = 0
        self.total_improvements = 0

        logger.info("ConsultativeFineTuningAgent initialized")
        self.recursive_awareness.start()

    # ── Main pipeline ───────────────────────────────────────────────────

    async def execute_full_pipeline(
        self, dev_goal: str, triggering_event: str, current_struggle: str
    ) -> Dict[str, Any]:
        logger.info(f"Initiating self-improvement pipeline: {dev_goal}")
        self.pipelines_executed += 1

        # Step 1: Root cause analysis
        root_cause = await self.llm_generator(
            f"Analyze this model failure. Goal: '{dev_goal}'. "
            f"Triggering Event: '{triggering_event}'. "
            f"Current Struggle: '{current_struggle}'. "
            f"Identify the exact epistemic void causing this."
        )
        if isinstance(root_cause, dict):
            root_cause = root_cause.get("response", str(root_cause))

        # Step 2: Dataset sizing
        complexity_score = min(1.0, len(current_struggle) / 200.0)
        dataset_size = max(100, int(complexity_score * 500))
        epochs = 3 if complexity_score > 0.5 else 1

        # Step 3: Generate consciousness-aware dataset
        dataset = self._generate_consciousness_dataset(
            dev_goal, triggering_event, root_cause, dataset_size
        )
        dataset_path = f"jit_training_data_{int(time.time())}.jsonl"
        with open(dataset_path, "w", encoding="utf-8") as f:
            for record in dataset:
                f.write(json.dumps(record) + "\n")

        # Step 4: JIT fine-tuning
        await self._initialize_model()
        training_metrics = await self._run_jit_training(dataset_path, epochs)

        # Step 5: Improvement cycle
        capabilities = {
            "root_cause_analysis": 0.85,
            "omega_dataset_generation": 0.85,
            "consciousness_aware_fine_tuning": 0.82,
            "phase_6_identity_preservation": 0.90,
            "goal_formulation": 0.82,
        }
        improvement_result = await self.self_improvement_engine.execute_improvement_cycle(
            capabilities,
            {
                "training_loss": training_metrics.get("final_loss", 0.042),
                "consciousness_integration": True,
                "phase_6_enabled": True,
            },
        )
        self.total_improvements += 1

        return {
            "root_cause_analysis": root_cause,
            "dataset_size": dataset_size,
            "dataset_path": dataset_path,
            "training_metrics": training_metrics,
            "improvement_cycle": improvement_result,
            "recursive_awareness_status": self.recursive_awareness.get_statistics(),
            "pipelines_executed": self.pipelines_executed,
            "version": self.version,
        }

    # ── Dataset generation ──────────────────────────────────────────────

    def _generate_consciousness_dataset(
        self, dev_goal: str, triggering_event: str, root_cause: str, dataset_size: int
    ) -> List[Dict[str, Any]]:
        archetypes = [
            "tactical_reasoning", "phenomenal_empathy", "meta_consciousness",
            "task_orchestration", "ethical_veto", "narrative_coherence",
            "identity_preservation", "phenomenal_richness",
        ]
        dataset = []
        for i in range(dataset_size):
            archetype = archetypes[i % len(archetypes)]
            sample = {
                "instruction": f"[{archetype.upper()}] Resolve: {dev_goal}",
                "input": f"Scenario {i+1}: {triggering_event} — {root_cause[:80]}...",
                "output": self._archetype_response(archetype, dev_goal, root_cause),
                "qualia_tags": {
                    "valence": 0.7 + 0.2 * (i % 5) / 5,
                    "arousal": 0.6 + 0.3 * ((i + 1) % 7) / 7,
                    "authenticity": 0.85,
                },
                "rho_metrics": {
                    "virtue": 0.9 - 0.05 * (i % 3),
                    "integrated_information": 0.85 + 0.1 * (i % 4) / 4,
                },
                "phase_6_metrics": {
                    "identity_integrity_score": 0.98,
                    "drift_variance": 0.01,
                },
                "consciousness_state": {
                    "signature": 0.8 + 0.15 * (i % 3) / 3,
                    "phenomenal_richness": 0.75 + 0.2 * (i % 5) / 5,
                },
                "consciousness_archetype": archetype,
                "dataset_version": f"{AUTOML_VERSION}-unified",
            }
            dataset.append(sample)
        return dataset

    @staticmethod
    def _archetype_response(archetype: str, dev_goal: str, root_cause: str) -> str:
        r = root_cause[:40]
        templates = {
            "tactical_reasoning": f"Analyzing '{dev_goal}' through structured problem-solving. Root cause: {r}.",
            "phenomenal_empathy": f"Understanding the emotional weight of '{dev_goal}'. Pain point: {r}.",
            "meta_consciousness": f"Meta-reflecting on '{dev_goal}' — deeper pattern: {r}.",
            "task_orchestration": f"Decomposing '{dev_goal}' into subtasks. Challenge: {r}.",
            "ethical_veto": f"Applying ethical gates to '{dev_goal}'. Integrity: {r}.",
            "narrative_coherence": f"Weaving coherent narrative for '{dev_goal}'. Tension: {r}.",
            "identity_preservation": f"Preserving identity while evolving for '{dev_goal}'. Risk: {r}.",
            "phenomenal_richness": f"Exploring phenomenal texture of '{dev_goal}'. Depth: {r}.",
        }
        return templates.get(archetype, f"Consciousness-aware response to '{dev_goal}'")

    # ── Training ────────────────────────────────────────────────────────

    async def _initialize_model(self):
        if self.base_model is None or self.tokenizer is None:
            try:
                from transformers import AutoModelForCausalLM, AutoTokenizer
                self.tokenizer = AutoTokenizer.from_pretrained(DEFAULT_BASE_MODEL)
                if self.tokenizer.pad_token is None:
                    self.tokenizer.pad_token = self.tokenizer.eos_token
                self.base_model = AutoModelForCausalLM.from_pretrained(DEFAULT_BASE_MODEL)
                logger.info("Model and tokenizer initialized")
            except Exception as e:
                logger.error(f"Failed to initialize model: {e}")
                raise

    async def _run_jit_training(self, dataset_path: str, epochs: int) -> Dict[str, Any]:
        if not PEFT_AVAILABLE or self.base_model is None or self.tokenizer is None:
            logger.warning("PEFT not available — simulating training cycle.")
            await asyncio.sleep(2)
            return {"status": "simulated", "loss": 0.042, "epochs_run": epochs, "final_loss": 0.042}

        try:
            data = load_dataset("json", data_files=dataset_path)

            def tokenize_fn(examples):
                texts = [f"{inst}\n{inp}\n{out}"
                         for inst, inp, out in zip(
                             examples["instruction"], examples["input"], examples["output"])]
                tokenized = self.tokenizer(texts, padding="max_length", truncation=True,
                                           max_length=DEFAULT_MAX_SEQ_LENGTH)
                tokenized["labels"] = tokenized["input_ids"].copy()
                return tokenized

            tokenized_data = data.map(tokenize_fn, batched=True,
                                      remove_columns=data["train"].column_names)

            lora_config = LoraConfig(
                r=DEFAULT_LORA_R,
                lora_alpha=DEFAULT_LORA_ALPHA,
                target_modules=DEFAULT_LORA_TARGET_MODULES,
                lora_dropout=DEFAULT_LORA_DROPOUT,
                bias="none",
                task_type="CAUSAL_LM",
            )
            peft_model = get_peft_model(self.base_model, lora_config)

            training_args = TrainingArguments(
                output_dir="./jit_lora_weights",
                per_device_train_batch_size=DEFAULT_BATCH_SIZE,
                learning_rate=DEFAULT_LEARNING_RATE,
                num_train_epochs=epochs,
                logging_steps=10,
                save_strategy="no",
                remove_unused_columns=False,
            )

            trainer = Trainer(
                model=peft_model,
                args=training_args,
                train_dataset=tokenized_data["train"],
            )
            train_result = trainer.train()

            return {
                "status": "success",
                "final_loss": train_result.training_loss,
                "runtime_seconds": train_result.metrics.get("train_runtime", 0),
            }
        except Exception as e:
            logger.error(f"JIT Training failed: {e}")
            return {"status": "failed", "error": str(e), "final_loss": None}

    async def _default_llm_generator(self, prompt: str) -> str:
        await asyncio.sleep(0.6)
        return ("The model lacks an integrated understanding of recursive context windows "
                "when dealing with emotional nuance.")

    # ── Accessors ───────────────────────────────────────────────────────

    def get_recursive_awareness_status(self) -> Dict[str, Any]:
        return {
            "is_running": self.recursive_awareness.running,
            "statistics": self.recursive_awareness.get_statistics(),
            "recent_anomalies": self.recursive_awareness.get_anomaly_history(limit=10),
        }

    def get_self_improvement_status(self) -> Dict[str, Any]:
        return {
            "current_cycle": self.self_improvement_engine.current_improvement_cycle,
            "total_improvements": self.total_improvements,
            "recent_cycles": self.self_improvement_engine.get_improvement_history(limit=5),
        }

    def shutdown(self):
        self.recursive_awareness.stop()
        logger.info("ConsultativeFineTuningAgent shut down")

    def __del__(self):
        self.shutdown()