api/speculative.py

"""
backend/api/speculative.py — Speculative Tool Firing (S361)

Pre-fires tool calls in parallel while the main model is processing.
Uses Groq llama-3.1-8b-instant for ultra-fast intent extraction (~200-300ms).
Results stored in a per-goal cache, consumed by _run_direct_tools before actual execution.

Architecture:
  1. Task created → fire_speculative_tools(task_id, goal) [fire-and-forget]
  2. Groq 8b extracts likely tools (JSON) in ~300ms
  3. Tools fired in parallel via TOOL_REGISTRY
  4. Results buffered in _spec_cache[goal_hash]
  5. _run_direct_tools checks cache via get_speculative_result() → cache hit = 0ms latency

Invariants:
  - NEVER blocks task creation or streaming response
  - NEVER raises — all errors are silent fallbacks
  - Cache TTL: 90s (expires before any reasonable task completes)
  - Only read-only, side-effect-free tools are speculated
"""
from __future__ import annotations

import asyncio
import hashlib
import json
import os
import re
import time
from typing import Any

# ── Speculative cache ─────────────────────────────────────────────────────────
# Keyed by goal_hash → {tool_key: result_str, _expires_at: float}
_spec_cache: dict[str, dict[str, Any]] = {}
_CACHE_TTL = 90.0   # seconds
_MAX_CACHE  = 50    # max concurrent entries (prevents unbounded growth)

# Only safe, read-only tools are eligible for speculation.
_SPECULATABLE: frozenset[str] = frozenset({
    "web_search", "get_weather", "get_news", "search_wikipedia",
    "get_currency", "fetch_url", "search_github", "get_location", "math_eval",
    # S-ORCH-8GAP FIX-GAP8: deep speculation — VFS read-only tools
    "read_file", "directory_tree", "file_search",
    # Added: git context + listing (added to _KEYWORD_SPEC_MAP, must be speculatable)
    "git_status", "git_diff", "list_files",
})

_EXTRACTION_PROMPT = (
    "Analyze this message and respond ONLY with a JSON array of tools to execute.\n"
    "Available tools:\n"
    "  web_search(query)      — search the web for information\n"
    "  get_weather(city)      — weather forecast for a city\n"
    "  get_news(query)        — latest news articles on a topic\n"
    "  search_wikipedia(query)— search Wikipedia for a topic\n"
    "  get_currency(from,to)  — currency/crypto exchange rate\n"
    "  fetch_url(url)         — fetch content from a URL\n"
    "  search_github(query)   — search GitHub repositories\n"
    "  get_location(query)    — geo lookup for a place\n"
    "  math_eval(expression)  — evaluate a math expression\n"
    "  read_file(path)        — read a file from the virtual filesystem\n"
    "  directory_tree(path)   — list directory structure\n"
    "  file_search(query)     — search for files by name or content\n"
    "  git_status()           — show current git status\n"
    "  git_diff()             — show recent git changes\n"
    "  list_files(path)       — list files in a directory\n"
    "Response format: [{\"tool\": \"name\", \"args\": {\"param\": \"value\"}}]\n"
    "If no tool is needed, respond: []\n"
    "Rules: be conservative — only suggest tools with >70% probability of being useful.\n\n"
    "Message: {message}"
)


# ── Cache helpers ─────────────────────────────────────────────────────────────

def _goal_hash(goal: str) -> str:
    return hashlib.md5(goal[:300].encode(), usedforsecurity=False).hexdigest()[:12]

def _tool_key(tool_name: str, args: dict) -> str:
    try:
        s = json.dumps(args, sort_keys=True, ensure_ascii=False)[:400]  # S608: 200→400
    except Exception:
        s = str(args)[:400]  # S608: 200→400
    return f"{tool_name}::{s}"

def get_speculative_result(goal: str, tool_name: str, args: dict) -> str | None:
    """
    Legge un risultato speculativo dalla cache.
    Ritorna None su cache miss o entry scaduta.
    """
    gh = _goal_hash(goal)
    bucket = _spec_cache.get(gh)
    if not bucket:
        return None
    if bucket.get("_expires_at", 0.0) < time.monotonic():
        _spec_cache.pop(gh, None)
        return None
    key = _tool_key(tool_name, args)
    return bucket.get(key)  # type: ignore[return-value]

def purge_speculative(goal: str) -> None:
    """Rimuove la cache per un goal (chiamato al termine del task)."""
    _spec_cache.pop(_goal_hash(goal), None)

def _prune_cache() -> None:
    now = time.monotonic()
    expired = [k for k, v in _spec_cache.items() if v.get("_expires_at", 0.0) < now]
    for k in expired:
        _spec_cache.pop(k, None)
    # Cap hard: rimuovi le entry più vecchie se sopra il limite
    if len(_spec_cache) > _MAX_CACHE:
        oldest = sorted(_spec_cache.keys(),
                        key=lambda k: _spec_cache[k].get("_expires_at", 0.0))
        for k in oldest[:len(_spec_cache) - _MAX_CACHE]:
            _spec_cache.pop(k, None)


# ── Tool extraction (Groq 8b fast path) ──────────────────────────────────────

# S388: Groq client singleton per speculative — evita new OpenAI() per ogni task.
_spec_groq_client: Any = None


def _get_spec_groq_client() -> Any:
    global _spec_groq_client
    if _spec_groq_client is not None:
        return _spec_groq_client
    groq_key = os.getenv("GROQ_API_KEY")
    if not groq_key:
        return None
    try:
        from openai import OpenAI
        _spec_groq_client = OpenAI(
            api_key=groq_key,
            base_url="https://api.groq.com/openai/v1",
            timeout=3.0,
            max_retries=0,
        )
    except Exception:
        _spec_groq_client = None
    return _spec_groq_client


async def _extract_tools_fast(goal: str) -> list[dict]:
    """
    Usa Groq llama-3.1-8b-instant per estrarre tool calls in ~300ms.
    Fallback silenzioso → [] se timeout, errore o key assente.
    """
    if not os.getenv("GROQ_API_KEY"):
        return []
    try:
        client = _get_spec_groq_client()
        if not client:
            return []
        prompt = _EXTRACTION_PROMPT.format(message=goal[:500])
        resp = await asyncio.wait_for(
            asyncio.to_thread(
                client.chat.completions.create,
                model="llama-3.1-8b-instant",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.0,
                max_tokens=400,  # S587: 256→400 — JSON array da goal[:500] supera 256 tok
            ),
            timeout=3.0,
        )
        raw = (resp.choices[0].message.content or "").strip()
        # Estrai JSON array anche se ci sono prefissi di testo
        start = raw.find("[")
        end   = raw.rfind("]") + 1
        if start == -1 or end == 0:
            return []
        parsed = json.loads(raw[start:end])
        if not isinstance(parsed, list):
            return []
        return [
            t for t in parsed
            if isinstance(t, dict)
            and isinstance(t.get("tool"), str)
            and t["tool"] in _SPECULATABLE
            and isinstance(t.get("args", {}), dict)
        ]
    except Exception:
        return []  # Fallback silenzioso totale


# ── Parallel tool executor ────────────────────────────────────────────────────

async def _run_tool_spec(tool_name: str, args: dict) -> tuple[str, str]:
    """
    Esegue un tool speculabile via TOOL_REGISTRY.
    Ritorna (cache_key, result_str) — result vuoto su errore.
    """
    key = _tool_key(tool_name, args)
    try:
        from tools.registry import TOOL_REGISTRY
        spec = TOOL_REGISTRY.get(tool_name)
        if not spec:
            return (key, "")
        fn = spec.get("_fn") or spec if callable(spec) else None
        if not fn:
            return (key, "")
        if asyncio.iscoroutinefunction(fn):
            result = await asyncio.wait_for(fn(**args), timeout=8.0)
        else:
            result = await asyncio.wait_for(asyncio.to_thread(fn, **args), timeout=8.0)
        return (key, str(result) if result is not None else "")
    except Exception:
        return (key, "")


# ── S-ORCH-8GAP FIX-GAP8: Deep File Speculation ─────────────────────────────

# Estrae path file espliciti menzionati nel goal (es. "correggi auth.py")
_FILE_PATH_RE = re.compile(
    r'\b([\w./\-]+\.(?:py|ts|tsx|js|jsx|css|html|json|ya?ml|sh|toml|sql|vue|svelte))\b',
    re.IGNORECASE,
)

# Keyword → tool da pre-eseguire per ridurre latenza planner
_KEYWORD_SPEC_MAP: list[tuple] = [
    # Auth/security → directory_tree (need file structure context)
    (re.compile(r'\b(login|auth|logout|session|token|jwt|oauth|password|signup|register|SSO|2fa|mfa)\b', re.I),
     {"tool": "directory_tree", "args": {"path": "."}}),
    # DB/ORM → directory_tree
    (re.compile(r'\b(database|schema|migration|model|orm|sql|db|postgres|mysql|sqlite|mongo|redis|supabase|prisma|drizzle)\b', re.I),
     {"tool": "directory_tree", "args": {"path": "."}}),
    # Bug/error → directory_tree + git_status (need recent changes context)
    (re.compile(r'\b(bug|fix|error|crash|exception|traceback|TypeError|ImportError|NameError|SyntaxError|broken|failing|not working)\b', re.I),
     {"tool": "git_status", "args": {}}),
    # API/routing → directory_tree
    (re.compile(r'\b(api|endpoint|route|handler|controller|middleware|fastapi|express|flask|django|router|webhook)\b', re.I),
     {"tool": "directory_tree", "args": {"path": "."}}),
    # Frontend/UI → directory_tree (need src structure)
    (re.compile(r'\b(frontend|component|react|vue|svelte|nextjs|vite|tailwind|css|styling|ui|design|layout|navbar|modal|button)\b', re.I),
     {"tool": "directory_tree", "args": {"path": "."}}),
    # DevOps/deploy → git_status (need current state)
    (re.compile(r'\b(deploy|deployment|ci|cd|github.?actions|dockerfile|docker.?compose|kubernetes|nginx|vercel|netlify|cloudflare)\b', re.I),
     {"tool": "git_status", "args": {}}),
    # Testing → directory_tree
    (re.compile(r'\b(test|testing|spec|jest|pytest|vitest|coverage|unit.?test|integration.?test|e2e)\b', re.I),
     {"tool": "directory_tree", "args": {"path": "."}}),
    # Package/deps → read_file package.json or requirements.txt
    (re.compile(r'\b(package|dependency|dependencies|install|npm|pip|pnpm|yarn|requirements|pyproject|setup.py)\b', re.I),
     {"tool": "read_file", "args": {"path": "package.json"}}),
    # Config/env → list_files
    (re.compile(r'\b(config|configuration|env|environment|secret|variable|\.env|settings|dotenv)\b', re.I),
     {"tool": "list_files", "args": {"path": "."}}),
    # Performance → git_diff (what changed recently?)
    (re.compile(r'\b(performance|slow|latency|optimize|speed|memory|cpu|profil|benchmark|bottleneck)\b', re.I),
     {"tool": "git_diff", "args": {}}),
]


async def _deep_speculate_files(goal: str, bucket: dict) -> None:
    """
    S-ORCH-8GAP FIX-GAP8: Pre-carica file VFS basandosi sul contenuto del goal.

    Step 1: estrai path file espliciti dal testo → read_file pre-caricato in cache
    Step 2: per keyword ad alto impatto → directory_tree (mappa struttura progetto)

    Risultati aggiunti al bucket → disponibili via get_speculative_result() prima
    che il planner/executor completino il routing. Riduce latenza percepita 2-5s.
    Fallback silenzioso totale — non blocca mai.
    """
    calls: list[dict] = []

    # Step 1: path file espliciti nel goal
    for m in _FILE_PATH_RE.finditer(goal[:500]):
        path = m.group(1)
        key  = _tool_key("read_file", {"path": path})
        if key not in bucket:
            calls.append({"tool": "read_file", "args": {"path": path}})

    # Step 2: keyword → directory_tree (max una sola per goal)
    _dir_key = _tool_key("directory_tree", {"path": "."})
    if _dir_key not in bucket:
        for pat, call in _KEYWORD_SPEC_MAP:
            if pat.search(goal[:400]):
                calls.append(call)
                break

    if not calls:
        return

    results = await asyncio.gather(
        *[_run_tool_spec(c["tool"], c.get("args", {})) for c in calls],
        return_exceptions=True,
    )
    for res in results:
        if isinstance(res, tuple) and res[1]:
            bucket[res[0]] = res[1]


# ── Entry point ───────────────────────────────────────────────────────────────

async def fire_speculative_tools(task_id: str, goal: str) -> None:
    """
    Chiamato fire-and-forget da create_agent_task().
    Non blocca mai, non lancia mai eccezioni.

    Fasi:
      1. Prune cache scaduta
      2. Estrai tool calls tramite Groq 8b (~300ms)
      3. Esegui tutti i tool in parallelo (max 8s each)
      4. Salva risultati nel buffer keyed by goal_hash
    """
    try:
        _prune_cache()

        # Skip se cache già calda per questo goal
        gh = _goal_hash(goal)
        if gh in _spec_cache:
            return

        # Fase 1: estrazione veloce
        tool_calls = await _extract_tools_fast(goal)
        if not tool_calls:
            return

        # Fase 2: esecuzione parallela
        results = await asyncio.gather(
            *[_run_tool_spec(tc["tool"], tc.get("args", {})) for tc in tool_calls],
            return_exceptions=True,
        )

        # Fase 3: scrivi in cache
        bucket: dict[str, Any] = {"_expires_at": time.monotonic() + _CACHE_TTL}
        for res in results:
            if isinstance(res, tuple) and res[1]:
                bucket[res[0]] = res[1]

        # S-ORCH-8GAP FIX-GAP8: deep speculation file VFS (after external tools)
        await _deep_speculate_files(goal, bucket)

        # Salva solo se almeno un tool ha restituito dati
        if len(bucket) > 1:  # > 1 perché _expires_at è sempre presente
            _spec_cache[gh] = bucket

    except Exception:
        pass  # Fallback silenzioso assoluto — non impatta mai il task