lambdagent/agentpaas/services/learning.py

"""
agentpaas.services.learning — Adaptive learning service

Records tool call patterns (success/failure), extracts strategy rules,
and auto-applies learned strategies in future similar scenarios.
"""
from __future__ import annotations

import json
import os
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List


@dataclass
class LearningEntry:
    tool_name: str
    input_preview: str
    success: bool
    error: str = ""
    strategy_used: str = ""     # e.g. "bash_python3_fallback"
    timestamp: float = 0

    def __post_init__(self):
        if self.timestamp == 0:
            self.timestamp = time.time()


@dataclass
class StrategyRule:
    pattern: str                # e.g. "WriteFile with markdown code blocks"
    strategy: str               # e.g. "Use Bash+python3 instead of WriteFile"
    confidence: float = 0.0     # 0.0 - 1.0
    evidence_count: int = 0
    last_applied: float = 0

    def to_dict(self):
        return {"pattern": self.pattern, "strategy": self.strategy,
                "confidence": self.confidence, "evidence_count": self.evidence_count}


LEARNING_PATH = os.path.expanduser("~/.lambdagent/learning.json")


class LearningService:
    """Track tool patterns and extract strategies."""

    def __init__(self, path: str = ""):
        self.path = path or LEARNING_PATH
        self._entries: List[LearningEntry] = []
        self._strategies: List[StrategyRule] = []
        self._load()

    def record(self, tool_name: str, input_preview: str, success: bool,
               error: str = "", strategy: str = ""):
        """Record a tool call outcome."""
        self._entries.append(LearningEntry(
            tool_name=tool_name, input_preview=input_preview[:200],
            success=success, error=error[:200], strategy_used=strategy,
        ))
        # Keep last 500 entries
        if len(self._entries) > 500:
            self._entries = self._entries[-500:]
        self._extract_rules()
        self._save()

    def get_strategy(self, tool_name: str, input_preview: str) -> str:
        """Get recommended strategy for a tool call."""
        for rule in sorted(self._strategies, key=lambda r: -r.confidence):
            if rule.confidence >= 0.6 and (
                tool_name.lower() in rule.pattern.lower() or
                any(w in input_preview.lower() for w in rule.pattern.lower().split())
            ):
                rule.last_applied = time.time()
                return rule.strategy
        return ""

    def list_strategies(self) -> List[StrategyRule]:
        return [r for r in self._strategies if r.confidence >= 0.5]

    def _extract_rules(self):
        """Auto-extract strategy rules from patterns."""
        # Count failures by tool
        fail_patterns: Dict[str, int] = {}
        success_patterns: Dict[str, int] = {}
        for e in self._entries:
            key = e.tool_name
            if e.success:
                success_patterns[key] = success_patterns.get(key, 0) + 1
            else:
                fail_patterns[key] = fail_patterns.get(key, 0) + 1

        # Strategy: if a tool fails >3 times and an alternative succeeds
        for tool, fails in fail_patterns.items():
            if fails >= 3:
                # Check if there's a successful alternative
                successes = success_patterns.get(tool, 0)
                total = fails + successes
                fail_rate = fails / max(1, total)
                if fail_rate > 0.5:
                    # Find what strategy succeeded
                    successful_strategies = [
                        e.strategy_used for e in self._entries
                        if e.tool_name == tool and e.success and e.strategy_used
                    ]
                    if successful_strategies:
                        strategy = max(set(successful_strategies), key=successful_strategies.count)
                        self._update_rule(
                            f"{tool} frequently fails",
                            f"Use alternative: {strategy}",
                            fail_rate,
                            fails,
                        )

    def _update_rule(self, pattern: str, strategy: str, confidence: float, evidence: int):
        for rule in self._strategies:
            if rule.pattern == pattern:
                rule.strategy = strategy
                rule.confidence = confidence
                rule.evidence_count = evidence
                return
        self._strategies.append(StrategyRule(
            pattern=pattern, strategy=strategy,
            confidence=confidence, evidence_count=evidence,
        ))

    def _save(self):
        os.makedirs(os.path.dirname(self.path), exist_ok=True)
        data = {
            "strategies": [r.to_dict() for r in self._strategies],
            "recent_entries": [
                {"tool": e.tool_name, "success": e.success, "error": e.error[:100],
                 "timestamp": e.timestamp}
                for e in self._entries[-50:]
            ],
        }
        with open(self.path, "w") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    def _load(self):
        if not os.path.exists(self.path):
            return
        try:
            with open(self.path) as f:
                data = json.load(f)
            for r in data.get("strategies", []):
                self._strategies.append(StrategyRule(**r))
        except Exception:
            pass


# Agent-side tool functions
_default_service = None

def _get_service() -> LearningService:
    global _default_service
    if _default_service is None:
        _default_service = LearningService()
    return _default_service


def learning_feedback(input_val: Any) -> str:
    params = _parse(input_val)
    _get_service().record(
        tool_name=params.get("tool", ""),
        input_preview=params.get("input", ""),
        success=params.get("success", True),
        error=params.get("error", ""),
        strategy=params.get("strategy", ""),
    )
    return "[OK] Feedback recorded"


def learning_strategies(input_val: Any) -> str:
    rules = _get_service().list_strategies()
    if not rules:
        return "[EMPTY] No strategies learned yet"
    lines = ["Learned strategies:"]
    for r in rules:
        lines.append(f"  [{r.confidence:.0%}] {r.pattern} → {r.strategy} ({r.evidence_count} evidence)")
    return "\n".join(lines)


def _parse(input_val: Any) -> dict:
    if isinstance(input_val, dict):
        return input_val
    if isinstance(input_val, str):
        try:
            return json.loads(input_val)
        except Exception:
            return {}
    return {}