reasoning_layer.py

# reasoning_layer.py – ADN NEXUS V11.2
# Boucles cognitives avancées : AST surgery, antagoniste, mémoire fractale, complexité dynamique

import ast
import re
import time
import json
import asyncio
from typing import Dict, Any, Optional, List, Tuple
from dataclasses import dataclass, field
import networkx as nx
import numpy as np
from sentence_transformers import SentenceTransformer
import astunparse  # pip install astunparse

# =================== 1. AST SURGEON ===================
class ASTSurgeon:
    """Répare le code en modifiant directement l'AST."""
    @staticmethod
    def repair_from_traceback(code: str, error_msg: str, line_no: int = 0) -> Tuple[str, bool]:
        try:
            tree = ast.parse(code)
            # Exemple : réparation d'erreur NameError
            if "NameError" in error_msg and "is not defined" in error_msg:
                match = re.search(r"name '(\w+)' is not defined", error_msg)
                if match:
                    var_name = match.group(1)
                    visitor = NameAdder(var_name, line_no)
                    visitor.visit(tree)
                    new_code = astunparse.unparse(tree)
                    return new_code, True
            # Autres corrections : import manquant, etc.
            return code, False
        except Exception as e:
            return code, False

class NameAdder(ast.NodeTransformer):
    def __init__(self, var_name, line_no):
        self.var_name = var_name
        self.line_no = line_no
    def visit_FunctionDef(self, node):
        if self.line_no == 0 or (node.lineno <= self.line_no <= node.end_lineno):
            new_stmt = ast.Assign(
                targets=[ast.Name(id=self.var_name, ctx=ast.Store())],
                value=ast.Constant(value=None)
            )
            node.body.insert(0, new_stmt)
            return node
        return node

# =================== 2. CENSEUR ANTAGONISTE ===================
class AntagonistCritic:
    def __init__(self, inference_client):
        self.inference = inference_client

    async def generate_tests(self, code: str, function_name: str = "main") -> List[str]:
        """Génère des cas de test limites en utilisant le modèle d'inférence."""
        prompt = f"""
Vous êtes un testeur impitoyable. Pour la fonction suivante :
{code}
Générez 3 cas de test (entrées) qui pourraient la faire échouer (valeurs extrêmes, types inattendus, etc.).
Répondez avec une liste Python simple, par exemple : [1, "test", None]
"""
        response = await self.inference.generate(
            [{"role": "user", "content": prompt}],
            max_tokens=150,
            temperature=0.7,
            do_sample=True
        )
        try:
            # Extraction de la liste
            match = re.search(r'\[.*\]', response, re.DOTALL)
            if match:
                tests = ast.literal_eval(match.group())
                return tests if isinstance(tests, list) else []
        except:
            pass
        return []

    async def validate_against_tests(self, code: str, tests: List[str]) -> bool:
        """Exécute le code avec chaque test et vérifie qu'il ne plante pas."""
        # Pour chaque test, on exécute le code avec la valeur en argument
        # (simplifié : on suppose que la fonction s'appelle 'main')
        for test in tests:
            # Construction d'un script d'appel
            test_code = code + f"\nresult = main({repr(test)})\nprint(result)"
            result = await self.sandbox.run_python(test_code)
            if not result.get("success", False):
                return False
        return True

# =================== 3. MÉMOIRE FRACTALE HOLOGRAPHIQUE ===================
class FractalMemory:
    def __init__(self):
        self.graph = nx.DiGraph()
        self.encoder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
        self.node_counter = 0

    def add_change(self, file_path: str, diff: str, semantic: str):
        """Ajoute un nœud de changement avec un résumé sémantique."""
        node_id = f"change_{self.node_counter}"
        self.node_counter += 1
        self.graph.add_node(
            node_id,
            type="change",
            file=file_path,
            diff=diff,
            summary=semantic,
            embedding=self.encoder.encode([semantic])[0].tolist()
        )
        # Lier aux nœuds précédents du même fichier
        prev = [n for n in self.graph.nodes if self.graph.nodes[n].get('file') == file_path and n != node_id]
        for p in prev:
            self.graph.add_edge(p, node_id, relation="follows")

    def retrieve_context(self, query: str, k: int = 5) -> List[Dict]:
        """Recherche les nœuds les plus pertinents par similarité vectorielle."""
        if len(self.graph.nodes) == 0:
            return []
        query_emb = self.encoder.encode([query])[0]
        nodes_with_emb = []
        for n in self.graph.nodes:
            emb = self.graph.nodes[n].get('embedding')
            if emb is not None:
                sim = np.dot(query_emb, emb) / (np.linalg.norm(query_emb) * np.linalg.norm(emb) + 1e-8)
                nodes_with_emb.append((n, sim))
        nodes_with_emb.sort(key=lambda x: x[1], reverse=True)
        top = nodes_with_emb[:k]
        return [self.graph.nodes[n] for n, _ in top]

# =================== 4. ESTIMATEUR DE COMPLEXITÉ ===================
class ComplexityEstimator:
    @staticmethod
    def estimate(task: str) -> float:
        """Retourne un score de complexité entre 0 et 1."""
        score = 0.0
        # Longueur
        if len(task) > 100:
            score += 0.3
        if len(task) > 200:
            score += 0.2
        # Mots-clés de complexité
        keywords = ["complex", "optimize", "refactor", "multi-thread", "performance", "algorithm", "design pattern"]
        if any(k in task.lower() for k in keywords):
            score += 0.4
        # Présence de code ou d'instructions spécifiques
        if "class" in task or "def " in task:
            score += 0.2
        return min(score, 1.0)

# =================== COUCHE DE RAISONNEMENT PRINCIPALE ===================
@dataclass
class CandidateResult:
    response: str
    tool_call: Optional[Dict]
    result: Dict
    score: float
    penalty: float = 0.0
    thought: Optional[str] = None
    execution_time_ms: float = 0.0

class ReasoningLayer:
    def __init__(self, inference_client, sandbox, tools, file_manager):
        self.inference = inference_client
        self.sandbox = sandbox
        self.tools = tools
        self.file_mgr = file_manager
        self.memory = {"successes": 0, "failures": 0, "recent_scores": []}
        # Nouvelles boucles
        self.ast_surgeon = ASTSurgeon()
        self.critic = AntagonistCritic(inference_client)
        self.fractal_memory = FractalMemory()
        self.complexity = ComplexityEstimator()

    @staticmethod
    def inject_cot_prompt(system_prompt: str) -> str:
        cot = """
BEFORE calling any tool, you MUST write a short reasoning block (max 50 words) in this format:
THOUGHT: <your step-by-step reasoning to solve the task>

Then, call the tool with a JSON object.

EXAMPLE:
User: "Run the test suite"
THOUGHT: The user wants to run tests. I need to execute pytest via run_python.
{"tool": "run_python", "arguments": {"code": "import subprocess\\nr = subprocess.run(['pytest', '-q'], capture_output=True, text=True)\\nprint(r.stdout)"}}"""
        return system_prompt + "\n\n" + cot

    @staticmethod
    def parse_thought(response: str) -> Optional[str]:
        match = re.search(r"THOUGHT:\s*(.*?)(?=\n|$)", response, re.DOTALL)
        return match.group(1).strip() if match else None

    @staticmethod
    def strip_thought(response: str) -> str:
        lines = response.split('\n')
        cleaned = [l for l in lines if not l.strip().startswith("THOUGHT:")]
        return '\n'.join(cleaned).strip()

    def evaluate_outcome(self, result: dict, execution_time_ms: float = 0) -> float:
        score = 0.0
        if result.get("success", False) and result.get("returncode", -1) == 0:
            score += 0.6
        elif result.get("success", False):
            score += 0.2
        if result.get("stderr", "") == "":
            score += 0.2
        stdout = result.get("stdout", "").strip()
        if stdout and len(stdout) > 10:
            score += 0.1
        if "error" in result.get("stderr", "").lower():
            score -= 0.1
        if execution_time_ms > 0 and execution_time_ms < 2000:
            score += 0.1
        return max(0.0, min(1.0, score))

    async def generate_with_penalty(self, messages, max_tokens=512, penalty_factor=0.05) -> Tuple[str, float]:
        response = await self.inference.generate(messages, max_tokens=max_tokens, do_sample=False)
        length = len(response.split())
        max_length = max_tokens // 2
        penalty = max(0.0, (length - max_length) / max_length) * penalty_factor
        return response, penalty

    async def run_with_sampling(self, task: str, websocket, max_candidates: int = None) -> Dict:
        # ==== BOUCLE 4 : COMPLEXITÉ DYNAMIQUE ====
        if max_candidates is None:
            complexity = self.complexity.estimate(task)
            max_candidates = max(2, int(2 + complexity * 3))  # 2 à 5 candidats
            max_tokens = int(512 * (1 + complexity * 0.5))
            await websocket.send_json({
                "type": "complexity_info",
                "complexity": complexity,
                "candidates": max_candidates,
                "max_tokens": max_tokens
            })
        else:
            max_tokens = 512

        # ==== BOUCLE 3 : MÉMOIRE FRACTALE ====
        context = self.fractal_memory.retrieve_context(task, k=3)
        context_text = ""
        if context:
            context_text = "\n[Contexte historique: " + "; ".join([c.get("summary", "") for c in context]) + "]"
            task_with_context = task + context_text
        else:
            task_with_context = task

        candidates = []
        base_messages = [{"role": "user", "content": task_with_context}]

        for i in range(max_candidates):
            prompt = f"{task_with_context} (approach {i+1}: {'direct and simple' if i == 0 else 'alternative method'})"
            start_time = time.time()
            response, penalty = await self.generate_with_penalty(
                base_messages + [{"role": "user", "content": prompt}],
                max_tokens=max_tokens,
                penalty_factor=0.05
            )
            execution_time = (time.time() - start_time) * 1000

            thought = self.parse_thought(response)
            cleaned_response = self.strip_thought(response)
            tool_call = self._parse_tool_call(cleaned_response)

            if tool_call:
                result = await self._exec_tool(tool_call["tool"], tool_call["arguments"])
                score = self.evaluate_outcome(result, execution_time) - penalty

                # ==== BOUCLE 1 : AST SURGERY en cas d'échec ====
                if not result.get("success", False) and tool_call.get("tool") == "run_python":
                    error = result.get("stderr", "") or result.get("error", "")
                    code = tool_call.get("arguments", {}).get("code", "")
                    if error and code:
                        repaired_code, ok = self.ast_surgeon.repair_from_traceback(code, error)
                        if ok:
                            # Réexécuter avec le code réparé
                            new_result = await self.sandbox.run_python(repaired_code)
                            if new_result.get("success", False):
                                result = new_result
                                score = self.evaluate_outcome(result, execution_time) - penalty

                # ==== BOUCLE 2 : VALIDATION ANTAGONISTE (optionnelle) ====
                if tool_call.get("tool") == "run_python" and result.get("success", False):
                    code = tool_call.get("arguments", {}).get("code", "")
                    if code:
                        tests = await self.critic.generate_tests(code)
                        if tests:
                            valid = await self.critic.validate_against_tests(code, tests)
                            if not valid:
                                score -= 0.2  # pénalité si les tests échouent
            else:
                result = {"success": True, "stdout": cleaned_response, "stderr": "", "returncode": 0}
                score = 0.5 - penalty

            candidates.append(CandidateResult(
                response=response, tool_call=tool_call, result=result,
                score=score, penalty=penalty, thought=thought,
                execution_time_ms=execution_time
            ))

            await websocket.send_json({
                "type": "sampling_candidate",
                "candidate": i + 1,
                "score": score,
                "thought": thought,
                "preview": result.get("stdout", "")[:200] if result else ""
            })

        # Sélection du meilleur candidat
        best = max(candidates, key=lambda x: x.score)
        self.memory["recent_scores"].append(best.score)
        if len(self.memory["recent_scores"]) > 10:
            self.memory["recent_scores"] = self.memory["recent_scores"][-10:]

        await websocket.send_json({
            "type": "sampling_summary",
            "best_candidate": candidates.index(best) + 1,
            "best_score": best.score,
            "scores": [c.score for c in candidates],
            "penalties": [c.penalty for c in candidates]
        })

        # Exécution finale du meilleur candidat
        if best.tool_call:
            final_result = await self._exec_tool(best.tool_call["tool"], best.tool_call["arguments"])
            # Stockage dans la mémoire fractale
            if best.tool_call.get("tool") == "run_python":
                self.fractal_memory.add_change(
                    file_path="generated.py",
                    diff=best.response,
                    semantic=f"Résolution de : {task}"
                )
        else:
            final_result = best.result

        return {
            "success": True,
            "result": final_result,
            "candidates_evaluated": len(candidates),
            "best_score": best.score,
            "thought_used": best.thought,
            "full_response": best.response
        }

    def _parse_tool_call(self, text: str) -> Optional[Dict]:
        decoder = json.JSONDecoder()
        start = 0
        while True:
            idx = text.find("{", start)
            if idx == -1:
                break
            try:
                obj, end = decoder.raw_decode(text, idx)
            except json.JSONDecodeError:
                start = idx + 1
                continue
            if isinstance(obj, dict) and "tool" in obj:
                return obj
            start = end
        return None

    async def _exec_tool(self, name: str, args: Dict) -> Dict:
        if name == "read_file":
            result = self.file_mgr.read(args.get("path", ""))
        elif name == "write_file":
            result = self.file_mgr.write(args.get("path", ""), args.get("content", ""))
        elif name == "list_directory":
            result = self.file_mgr.list_dir(args.get("path", ""))
        elif name == "run_python":
            result = await self.sandbox.run_python(args.get("code", ""))
        else:
            result = {"success": False, "error": f"Tool {name} unknown"}
        return result

    def get_stats(self) -> Dict:
        avg_score = sum(self.memory["recent_scores"]) / max(len(self.memory["recent_scores"]), 1)
        return {
            "successes": self.memory["successes"],
            "failures": self.memory["failures"],
            "avg_recent_score": round(avg_score, 3),
            "total_evaluations": len(self.memory["recent_scores"])
        }