goal_engine.py

#!/usr/bin/env python3.12
"""ULTRON Goal Engine - autonomous goal tracking with verification and learning."""

import argparse
import json
import logging
import os
import shutil
from datetime import datetime, timezone
from pathlib import Path

import requests
import yaml

from actions import parse_actions, execute_parsed_action
from learnings import (
    record_failure, record_success, get_anti_patterns_for_task,
    get_successful_approaches, get_stats as get_learning_stats,
)
from verify import Verifier, verify_task_completion, should_ask_human

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [goal_engine] %(levelname)s %(message)s",
)
log = logging.getLogger(__name__)

BASE_DIR = Path(__file__).resolve().parent
GOALS_PATH = BASE_DIR / "goals.yaml"
RESULTS_DIR = BASE_DIR / "results"
PRIORITY_MAP = {"high": 3, "medium": 2, "low": 1}
MAX_RETRIES = 3


def load_config():
    with open(BASE_DIR / "config.yaml") as f:
        return yaml.safe_load(f)


def load_goals():
    if not GOALS_PATH.exists():
        log.warning("goals.yaml not found, creating with empty goals")
        GOALS_PATH.write_text("goals: []\n")
        return {"goals": []}
    with open(GOALS_PATH) as f:
        return yaml.safe_load(f) or {"goals": []}


def save_goals(data):
    if GOALS_PATH.exists():
        shutil.copy2(GOALS_PATH, GOALS_PATH.with_suffix(".yaml.bak"))
    with open(GOALS_PATH, "w") as f:
        yaml.dump(data, f, default_flow_style=False, sort_keys=False, width=120)


def ask_ollama(cfg, prompt):
    url = cfg["ollama"]["url"]
    for model in [cfg["ollama"]["model"], cfg["ollama"]["fallback_model"]]:
        try:
            resp = requests.post(url, json={
                "model": model,
                "prompt": prompt,
                "stream": False,
                "options": {"temperature": 0.5},
            }, timeout=180)
            resp.raise_for_status()
            return resp.json().get("response", "")
        except Exception as e:
            log.warning("Ollama failed (model=%s): %s", model, e)
    return None


def send_telegram(text):
    token = os.environ.get("TELEGRAM_BOT_TOKEN", "")
    chat_id = os.environ.get("TELEGRAM_CHAT_ID", "")
    if not token or not chat_id:
        log.info("Telegram not configured, skipping notification")
        return
    try:
        requests.post(
            f"https://api.telegram.org/bot{token}/sendMessage",
            json={"chat_id": chat_id, "text": text, "parse_mode": "Markdown"},
            timeout=15,
        )
    except Exception as e:
        log.warning("Telegram send failed: %s", e)


def store_in_aoms(cfg, goal, task, result_preview, verified=False):
    body = {
        "type": "experience",
        "payload": {
            "title": f"Task completed: {task['name']}",
            "goal": goal["name"],
            "outcome": "success",
            "summary": result_preview[:200],
            "verified": verified,
        },
        "tags": ["goal", goal["id"]],
    }
    try:
        requests.post(cfg["aoms"]["url"] + "/memory/episodic", json=body, timeout=10)
    except Exception as e:
        log.warning("AOMS store failed: %s", e)


def calc_progress(goal):
    tasks = goal.get("tasks", [])
    if not tasks:
        return 0
    done = sum(1 for t in tasks if t["status"] == "completed")
    return int((done / len(tasks)) * 100)


def find_next_task(goals_data):
    """Find highest-priority goal's next ready task."""
    active = [g for g in goals_data["goals"] if g.get("status") == "active"]
    active.sort(key=lambda g: PRIORITY_MAP.get(g.get("priority", "low"), 0), reverse=True)

    for goal in active:
        completed_ids = {t["id"] for t in goal.get("tasks", []) if t["status"] == "completed"}
        for task in goal.get("tasks", []):
            if task["status"] not in ("pending", "retry"):
                continue
            deps = set(task.get("requires", []))
            if deps.issubset(completed_ids):
                return goal, task
    return None, None


def build_task_prompt(goal, task):
    """Build task prompt with anti-patterns and successful approaches."""
    avoid = get_anti_patterns_for_task(task["name"])
    successes = get_successful_approaches(task["name"])

    context_parts = []
    if avoid:
        context_parts.append(f"\nAVOID THESE APPROACHES (failed before): {avoid}")
    if successes:
        context_parts.append(f"\nAPPROACHES THAT WORKED FOR SIMILAR TASKS: {successes}")

    retry_count = task.get("retry_count", 0)
    retry_note = ""
    if retry_count > 0:
        retry_note = f"\n\nNOTE: This is attempt {retry_count + 1}. Previous attempts failed verification. Try a DIFFERENT approach."

    return f"""You are ULTRON, an autonomous AI agent.

GOAL: {goal['name']}
TASK: {task['name']}{retry_note}

You can take these actions by including directives in your response:
- ACTION:TWEET text="your tweet text"
- ACTION:WRITE_FILE path="outputs/filename.md" content="file content here"
- ACTION:NOTIFY message="notification message"
- ACTION:GITHUB_CREATE_REPO name="repo-name" description="repo description"

Think through the task step by step. Provide a comprehensive result.
For research tasks: synthesize findings with sources/reasoning.
For writing tasks: produce the actual content and use ACTION:WRITE_FILE to save it.
For planning tasks: create actionable steps.

Be thorough. This result will be stored in memory and reviewed.{"".join(context_parts)}"""


def execute_task(cfg, goal, task):
    """Execute a task — uses deep_think for research tasks, single-shot for action tasks."""
    # Research-type tasks benefit from deep thinking
    research_keywords = ["research", "analyze", "investigate", "compare", "study", "explore", "understand"]
    is_research = any(kw in task["name"].lower() for kw in research_keywords)

    if is_research:
        try:
            from deep_think import deep_think
            log.info("Using DEEP THINK for research task")
            result = deep_think(
                f"Goal: {goal['name']}. Task: {task['name']}",
                use_web=True,
                use_code=True,
                cfg=cfg,
            )
            if result and result.get("final_answer"):
                return result["final_answer"]
            log.warning("Deep think failed for task, falling back to single-shot")
        except Exception as e:
            log.warning("Deep think error: %s, falling back", e)

    prompt = build_task_prompt(goal, task)
    return ask_ollama(cfg, prompt)


def save_result(goal, task, result):
    result_dir = RESULTS_DIR / goal["id"]
    result_dir.mkdir(parents=True, exist_ok=True)
    result_path = result_dir / f"{task['id']}.md"
    header = f"# {task['name']}\n\n**Goal:** {goal['name']}\n**Completed:** {datetime.now(timezone.utc).isoformat()}\n\n---\n\n"
    result_path.write_text(header + result)
    return result_path


def log_cycle(goal, task, success, verification=None):
    log_dir = BASE_DIR / "logs"
    log_dir.mkdir(exist_ok=True)
    entry = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "goal_id": goal["id"],
        "task_id": task["id"],
        "task_name": task["name"],
        "success": success,
        "retry_count": task.get("retry_count", 0),
    }
    if verification:
        entry["verification"] = {
            "method": verification.method,
            "evidence": verification.evidence,
            "confidence": verification.confidence,
        }
    with open(log_dir / "goal_engine.log", "a") as f:
        f.write(json.dumps(entry) + "\n")


def progress_bar(pct, width=16):
    filled = int(width * pct / 100)
    return "\u2588" * filled + "\u2591" * (width - filled)


def days_until(deadline_str):
    try:
        dl = datetime.strptime(deadline_str, "%Y-%m-%d").replace(tzinfo=timezone.utc)
        delta = dl - datetime.now(timezone.utc)
        return delta.days
    except (ValueError, TypeError):
        return None


def cmd_list(goals_data):
    print("\nULTRON Goal Tracker")
    print("\u2550" * 40)

    for goal in goals_data.get("goals", []):
        pct = calc_progress(goal)
        total = len(goal.get("tasks", []))
        done = sum(1 for t in goal.get("tasks", []) if t["status"] == "completed")
        failed = sum(1 for t in goal.get("tasks", []) if t["status"] == "failed")
        pri = goal.get("priority", "low").upper()
        pri_label = {"HIGH": "HIGH", "MEDIUM": "MED ", "LOW": "LOW "}.get(pri, pri)
        days = days_until(goal.get("deadline"))
        status = goal.get("status", "active")

        print(f"\n[{pri_label}] {goal['name']}")
        status_str = f"{pct}% ({done}/{total} tasks)"
        if failed:
            status_str += f", {failed} failed"
        print(f"       {progress_bar(pct)} {status_str}")
        if days is not None:
            print(f"       Deadline: {goal.get('deadline')} ({days} days)")
        if status not in ("active",):
            print(f"       Status: {status}")

        if status == "active":
            completed_ids = {t["id"] for t in goal.get("tasks", []) if t["status"] == "completed"}
            for task in goal.get("tasks", []):
                if task["status"] in ("pending", "retry") and set(task.get("requires", [])).issubset(completed_ids):
                    retry = task.get("retry_count", 0)
                    suffix = f" (retry #{retry})" if retry else ""
                    print(f"       Next: {task['name']}{suffix}")
                    break
                elif task["status"] == "awaiting_human":
                    print(f"       BLOCKED: {task['name']} (awaiting human input)")
                    break
    print()


def cmd_status(goals_data):
    active = [g for g in goals_data.get("goals", []) if g.get("status") == "active"]
    total_tasks = sum(len(g.get("tasks", [])) for g in active)
    done_tasks = sum(1 for g in active for t in g.get("tasks", []) if t["status"] == "completed")
    failed_tasks = sum(1 for g in active for t in g.get("tasks", []) if t["status"] == "failed")
    print(f"Goals: {len(active)} active | Tasks: {done_tasks}/{total_tasks} completed, {failed_tasks} failed")
    goal, task = find_next_task(goals_data)
    if goal and task:
        print(f"Next: [{goal.get('priority', '?').upper()}] {task['name']}")


def cmd_learnings():
    stats = get_learning_stats()
    print("\nULTRON Learning Stats")
    print("-" * 30)
    print(f"Successes: {stats['total_successes']}")
    print(f"Failures:  {stats['total_failures']}")
    print(f"Anti-patterns: {stats['anti_patterns']}")
    print(f"Unique failed tasks: {stats['unique_failed_tasks']}")


def main():
    parser = argparse.ArgumentParser(description="ULTRON Goal Engine")
    parser.add_argument("--dry-run", action="store_true", help="Show what would execute without doing it")
    parser.add_argument("--list", action="store_true", help="Show all goals with progress")
    parser.add_argument("--status", action="store_true", help="Quick status summary")
    parser.add_argument("--learnings", action="store_true", help="Show failures and anti-patterns")
    args = parser.parse_args()

    goals_data = load_goals()

    if args.list:
        cmd_list(goals_data)
        return
    if args.status:
        cmd_status(goals_data)
        return
    if args.learnings:
        cmd_learnings()
        return

    cfg = load_config()
    goal, task = find_next_task(goals_data)

    if not goal:
        log.info("No tasks ready to execute. All done or blocked.")
        return

    log.info("Goal: %s", goal["name"])
    log.info("Task: %s (attempt %d)", task["name"], task.get("retry_count", 0) + 1)

    if args.dry_run:
        log.info("[DRY RUN] Would execute task: %s", task["name"])
        return

    # Execute
    result = execute_task(cfg, goal, task)
    if not result:
        log.error("Task execution failed (Ollama unreachable)")
        record_failure(task["name"], "ollama_reasoning", "Ollama unreachable", goal["name"])
        log_cycle(goal, task, success=False)
        return

    # Execute ACTION directives
    parsed = parse_actions(result)
    if parsed:
        log.info("Found %d action(s) in response", len(parsed))
        for action_type, params in parsed:
            log.info("Executing: ACTION:%s %s", action_type, {k: v[:50] for k, v in params.items()})
            action_result = execute_parsed_action(action_type, params)
            log.info("Action result: %s", action_result)

    # Save result
    task["_goal_id"] = goal["id"]
    result_path = save_result(goal, task, result)
    log.info("Result saved: %s", result_path)

    # Verify
    verification = verify_task_completion(task, result)
    log.info("Verification: %s (method=%s, confidence=%.1f, evidence=%s)",
             "PASS" if verification.success else "FAIL",
             verification.method, verification.confidence, verification.evidence)

    if verification.success:
        # Check if human approval needed for high-risk tasks
        if should_ask_human(verification, task):
            log.info("Human approval required for this task")
            task["status"] = "awaiting_human"
            send_telegram(f"Need your input:\n\nTask: {task['name']}\nVerification: {verification.evidence}\n\nReply YES to approve.")
            save_goals(goals_data)
            log_cycle(goal, task, success=True, verification=verification)
            return

        record_success(task["name"], "ollama_reasoning", result[:500])
        task["status"] = "completed"
        goal["progress"] = calc_progress(goal)
        save_goals(goals_data)
        log.info("Progress: %d%%", goal["progress"])

        store_in_aoms(cfg, goal, task, result, verified=True)
        log_cycle(goal, task, success=True, verification=verification)

        total = len(goal.get("tasks", []))
        done = sum(1 for t in goal.get("tasks", []) if t["status"] == "completed")
        msg = (f"Verified Task: {task['name']}\n"
               f"Goal: {goal['name']}\n"
               f"Progress: {goal['progress']}% ({done}/{total} tasks)\n"
               f"Proof: {verification.evidence}")

        if goal["progress"] == 100:
            msg += "\n\nGoal COMPLETED!"
            goal["status"] = "completed"
            save_goals(goals_data)

        send_telegram(msg)
    else:
        # Verification failed
        retry_count = task.get("retry_count", 0) + 1
        task["retry_count"] = retry_count
        record_failure(task["name"], "ollama_reasoning", verification.evidence, goal["name"])

        if retry_count < MAX_RETRIES:
            task["status"] = "retry"
            save_goals(goals_data)
            log.warning("Verification failed, will retry (attempt %d/%d)", retry_count, MAX_RETRIES)
            send_telegram(f"Task failed verification (attempt {retry_count}/{MAX_RETRIES}): {task['name']}\nReason: {verification.evidence}")
        else:
            # Max retries exhausted — try adaptive replanning instead of giving up
            log.warning("Max retries exhausted, attempting adaptive replan...")
            try:
                from adaptive_planner import replan_failed_task, apply_replan
                plan = replan_failed_task(cfg, goal, task, verification.evidence)
                if plan and plan.get("action") in ("decompose", "alternative"):
                    apply_replan(goals_data, goal["id"], task["id"], plan)
                    new_tasks = [t["name"] for t in plan.get("tasks", [])]
                    log.info("Replanned: %s -> %s", task["name"], new_tasks)
                    send_telegram(
                        f"Task replanned: {task['name']}\n"
                        f"Action: {plan['action']}\n"
                        f"New tasks: {', '.join(new_tasks)}"
                    )
                else:
                    task["status"] = "failed"
                    save_goals(goals_data)
                    log.error("Replan failed, task marked as failed: %s", task["name"])
                    send_telegram(f"Task FAILED (replan also failed): {task['name']}\nReason: {verification.evidence}")
            except Exception as e:
                log.error("Adaptive replan error: %s", e)
                task["status"] = "failed"
                save_goals(goals_data)
                send_telegram(f"Task FAILED after {MAX_RETRIES} attempts: {task['name']}\nReason: {verification.evidence}")

        log_cycle(goal, task, success=False, verification=verification)

    log.info("Cycle complete.")


if __name__ == "__main__":
    main()