macrocosmos_agent.py

"""
Agente Macrocosmos: recibe instrucciones en lenguaje natural, interpreta con Chutes
y ejecuta las herramientas MCP (OnDemand + Gravity) para obtener datos.
"""

from __future__ import annotations

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

import requests
from dotenv import load_dotenv

import macrocosmos_tools as mt

load_dotenv()

CHUTES_API_TOKEN = os.getenv("CHUTES_API_TOKEN")
CHUTES_URL = "https://llm.chutes.ai/v1/chat/completions"

TOOLS_SCHEMA = """
Herramientas Macrocosmos MCP disponibles:

1. query_on_demand_data - Consultas en tiempo real (X o Reddit). Hasta 1000 resultados.
   Params: source (string, "X" o "REDDIT"), usernames (lista hasta 5, solo X), keywords (lista hasta 5; Reddit: primero = subreddit ej. r/MachineLearning), start_date (ISO opcional), end_date (ISO opcional), limit (1-1000, default 100), keyword_mode ("any" o "all", default "any").
   Ejemplos: "tweets de @elonmusk", "posts de r/bittensor sobre dTAO", "50 tweets sobre #AI la última semana".

2. create_gravity_task - Crear tarea de recolección masiva (7 días). Para más de 1000 resultados.
   Params: tasks (lista de objetos {"platform": "x" o "reddit", "topic": "#Bittensor" o "r/MachineLearning"; X debe empezar por # o $), name (opcional), email (opcional).
   Ejemplos: "crear tarea gravity para #Bittensor en X", "empezar a recolectar de r/MachineLearning sobre neural networks".

3. get_gravity_task_status - Ver progreso de una tarea Gravity.
   Params: gravity_task_id (string), include_crawlers (bool, default true).

4. build_dataset - Construir dataset desde un crawler (detiene el crawler).
   Params: crawler_id (string), max_rows (int, default 10000), email (opcional).

5. get_dataset_status - Estado del build y URLs de descarga.
   Params: dataset_id (string).

6. cancel_gravity_task - Cancelar tarea Gravity.
   Params: gravity_task_id (string).

7. cancel_dataset - Cancelar build o purgar dataset.
   Params: dataset_id (string).

Interpreta la instrucción del usuario y elige UNA herramienta con los parámetros apropiados. Para fechas relativas ("última semana", "hoy") usa ISO UTC. Reddit: subreddit como primer keyword (r/...). X: hashtags como #AI, símbolos como $BTC.
"""


def _extract_json(raw: str) -> dict[str, Any]:
    """Extrae un objeto JSON del texto (permite markdown code blocks)."""
    raw = raw.strip()
    # Quitar bloques ```json ... ```
    m = re.search(r"```(?:json)?\s*([\s\S]*?)```", raw)
    if m:
        raw = m.group(1).strip()
    # Buscar primer { ... } bien balanceado
    start = raw.find("{")
    if start < 0:
        raise ValueError("No se encontró ningún objeto JSON en la respuesta")
    depth = 0
    end = -1
    in_str = False
    escape = False
    quote = None
    i = start
    while i < len(raw):
        c = raw[i]
        if escape:
            escape = False
            i += 1
            continue
        if c == '\\' and in_str:
            escape = True
            i += 1
            continue
        if not in_str:
            if c in '"\'':
                in_str = True
                quote = c
            elif c == "{":
                depth += 1
            elif c == "}":
                depth -= 1
                if depth == 0:
                    end = i
                    break
        else:
            if c == quote:
                in_str = False
        i += 1
    if end < 0:
        raise ValueError("JSON malformado o no cerrado")
    return json.loads(raw[start : end + 1])


def _chutes_request(prompt: str) -> str:
    """Envía prompt a Chutes y devuelve el contenido de la respuesta."""
    if not CHUTES_API_TOKEN:
        raise RuntimeError("CHUTES_API_TOKEN no configurado en .env")
    headers = {
        "Authorization": f"Bearer {CHUTES_API_TOKEN}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": "Qwen/Qwen3-32B",
        "messages": [
            {
                "role": "system",
                "content": "Eres un asistente que traduce instrucciones de usuario a llamadas de la API Macrocosmos MCP. Debes responder ÚNICAMENTE con un objeto JSON válido, sin texto adicional ni explicaciones. Formato: {\"tool\": \"nombre_herramienta\", \"params\": {...}}.",
            },
            {
                "role": "user",
                "content": f"{TOOLS_SCHEMA}\n\nInstrucción del usuario:\n{prompt}\n\nResponde solo con el JSON (tool + params).",
            },
        ],
        "stream": False,
        "max_tokens": 1024,
        "temperature": 0.1,
    }
    r = requests.post(CHUTES_URL, headers=headers, json=payload, timeout=120)
    r.raise_for_status()
    data = r.json()
    if not data.get("choices"):
        raise RuntimeError("Respuesta vacía de Chutes")
    msg = data["choices"][0].get("message", {})
    return (msg.get("content") or "").strip()


def interpret_instruction(instruction: str) -> dict[str, Any]:
    """
    Interpreta la instrucción con Chutes y devuelve {tool, params}.
    """
    raw = _chutes_request(instruction)
    out = _extract_json(raw)
    tool = out.get("tool")
    params = out.get("params") or {}
    if not tool or not isinstance(tool, str):
        raise ValueError("El JSON debe incluir 'tool' (string)")
    if not isinstance(params, dict):
        raise ValueError("'params' debe ser un objeto")
    return {"tool": tool.strip(), "params": params}


async def _execute(tool: str, params: dict[str, Any]) -> Any:
    """Ejecuta la herramienta Macrocosmos indicada."""
    t = tool.strip().lower()
    p = {k: v for k, v in params.items() if v is not None}

    if t == "query_on_demand_data":
        return await mt.query_on_demand_data(
            source=p.get("source", "X"),
            usernames=p.get("usernames"),
            keywords=p.get("keywords"),
            start_date=p.get("start_date"),
            end_date=p.get("end_date"),
            limit=int(p.get("limit", 100)),
            keyword_mode=p.get("keyword_mode", "any"),
        )
    if t == "create_gravity_task":
        tasks = p.get("tasks")
        if not tasks or not isinstance(tasks, list):
            raise ValueError("create_gravity_task requiere 'tasks' (lista)")
        return await mt.create_gravity_task(
            tasks=tasks,
            name=p.get("name"),
            email=p.get("email"),
        )
    if t == "get_gravity_task_status":
        gid = p.get("gravity_task_id")
        if not gid:
            raise ValueError("get_gravity_task_status requiere 'gravity_task_id'")
        return await mt.get_gravity_task_status(
            gravity_task_id=str(gid),
            include_crawlers=p.get("include_crawlers", True),
        )
    if t == "build_dataset":
        cid = p.get("crawler_id")
        if not cid:
            raise ValueError("build_dataset requiere 'crawler_id'")
        return await mt.build_dataset(
            crawler_id=str(cid),
            max_rows=int(p.get("max_rows", 10000)),
            email=p.get("email"),
        )
    if t == "get_dataset_status":
        did = p.get("dataset_id")
        if not did:
            raise ValueError("get_dataset_status requiere 'dataset_id'")
        return await mt.get_dataset_status(dataset_id=str(did))
    if t == "cancel_gravity_task":
        gid = p.get("gravity_task_id")
        if not gid:
            raise ValueError("cancel_gravity_task requiere 'gravity_task_id'")
        return await mt.cancel_gravity_task(gravity_task_id=str(gid))
    if t == "cancel_dataset":
        did = p.get("dataset_id")
        if not did:
            raise ValueError("cancel_dataset requiere 'dataset_id'")
        return await mt.cancel_dataset(dataset_id=str(did))

    raise ValueError(f"Herramienta desconocida: {tool}")


def _format_result(result: Any, max_items: int = 50) -> str:
    """Formatea el resultado para mostrarlo al usuario."""
    if isinstance(result, list):
        lines = [f"Total: {len(result)} resultados."]
        for i, item in enumerate(result[:max_items]):
            if isinstance(item, dict):
                txt = item.get("texto") or item.get("content") or str(item)[:200]
                author = item.get("autor") or item.get("author") or "?"
                ts = item.get("timestamp") or item.get("datetime") or ""
                lines.append(f"  [{i+1}] {author} | {ts}")
                lines.append(f"      {txt[:300]}{'...' if len(str(txt)) > 300 else ''}")
            else:
                lines.append(f"  [{i+1}] {item}")
        if len(result) > max_items:
            lines.append(f"  ... y {len(result) - max_items} más.")
        return "\n".join(lines)
    if isinstance(result, dict):
        return json.dumps(result, indent=2, ensure_ascii=False)
    return str(result)


def run(instruction: str) -> str:
    """
    Interpreta la instrucción, ejecuta la herramienta Macrocosmos y devuelve
    el resultado formateado.
    """
    try:
        parsed = interpret_instruction(instruction)
    except (ValueError, RuntimeError, requests.RequestException) as e:
        return f"[Error interpretando] {e}"

    tool = parsed["tool"]
    params = parsed["params"]

    try:
        result = asyncio.run(_execute(tool, params))
    except Exception as e:
        return f"[Error ejecutando {tool}] {e}"

    return _format_result(result)