Lesson 05 — Cognee + autonomía acotada

"""Query the local Cognee graph.

⚠️ Ver `setup.md` sección "Bugs de integración con Cognee" para el contexto de
los mappings de env vars y del override post-import. Los patrones no son
opcionales — sin ellos falla con LLMAPIKeyNotSetError o 404 de OpenRouter.
"""

from __future__ import annotations

import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


# ─── Env var mapping (bug 2 en setup.md) ──────────────────────────
# cognee 1.4 lee LLM_* / EMBEDDING_* SIN el prefijo COGNEE_.
# Nuestro .env usa COGNEE_LLM_* por claridad; mapeamos acá al importar.
_COGNEE_ENV_VARS = (
    "LLM_API_KEY", "LLM_MODEL", "LLM_PROVIDER", "LLM_ENDPOINT",
    "EMBEDDING_API_KEY", "EMBEDDING_MODEL", "EMBEDDING_PROVIDER",
    "EMBEDDING_ENDPOINT", "EMBEDDING_DIMENSIONS",
)


def _map_cognee_env() -> None:
    """Copia COGNEE_<VAR> → <VAR> para que cognee 1.4 las encuentre."""
    for var in _COGNEE_ENV_VARS:
        val = os.environ.get(f"COGNEE_{var}") or os.environ.get(var)
        if val:
            os.environ[var] = val


_map_cognee_env()  # antes de import cognee

import cognee  # noqa: E402  ← import después del mapping intencional
from cognee.api.v1.search import SearchType  # noqa: E402

_map_cognee_env()  # DESPUÉS también: cognee re-carga .env con override=True

from src.agent import write_receipt  # noqa: E402


async def ingest(paths: list[Path]) -> None:
    for p in paths:
        await cognee.add(p.read_text(encoding="utf-8"), dataset_name="acme")
    await cognee.cognify(datasets=["acme"])
    write_receipt("cognee_indexed", {"files": len(paths), "dataset": "acme"})


async def query(question: str) -> dict[str, Any]:
    results = await cognee.search(
        query_text=question,
        query_type=SearchType.INSIGHTS,
        datasets=["acme"],
    )
    if not results:
        return {"answer": "no data", "source": "none", "confidence": 0.0, "raw": []}
    # Cognee INSIGHTS devuelve una lista de tuplas/dicts según versión.
    first = results[0]
    return {
        "answer": str(first),
        "source": "acme dataset",
        "confidence": 0.8,
        "raw": results[:3],
    }

"""DeepAgent lifecycle — explicit states over hidden booleans."""

from __future__ import annotations

import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from statemachine import State, StateMachine

from src.schemas import Receipt


class DeepAgentLifecycle(StateMachine):
    raw = State(initial=True)
    learning = State()
    answering = State()
    refining = State()
    escalated = State(final=True)

    detect_gap = raw.to(learning) | answering.to(learning)
    skill_ready = learning.to(answering)
    verify_ok = answering.to(answering)
    verify_fail = answering.to(refining)
    retry = refining.to(answering)
    give_up = refining.to(escalated) | learning.to(escalated)

    def on_transition(self, event: str, source: State, target: State) -> None:
        _write(event=event, payload={"from": source.id, "to": target.id})


def _write(event: str, payload: dict[str, Any]) -> None:
    receipts_dir = Path(os.environ.get("RECEIPTS_DIR", ".receipts"))
    receipts_dir.mkdir(exist_ok=True)
    path = receipts_dir / "lifecycle.jsonl"
    receipt = Receipt(
        ts=datetime.now(tz=timezone.utc),
        actor="system",
        event="transition",
        payload={"trigger": event, **payload},
    )
    with path.open("a", encoding="utf-8") as f:
        f.write(receipt.model_dump_json() + "\n")

# nuevo comando async
import asyncio
from pathlib import Path

from src.lifecycle import DeepAgentLifecycle
from src.skill_registry import discover_skills


@app.command()
def ask(question: str) -> None:
    asyncio.run(_ask_async(question))


async def _ask_async(question: str) -> None:
    lc = DeepAgentLifecycle()
    skills = discover_skills(Path(".skills"))

    if not any(s.name == "query_cognee" for s in skills):
        lc.detect_gap()   # raw → learning
        # ... invocar skill_author como antes ...
        lc.skill_ready()  # learning → answering
    else:
        lc.detect_gap()
        lc.skill_ready()

    # Ingesta idempotente (Cognee ignora duplicados con el mismo hash)
    from importlib import import_module
    tool = import_module("skills.query_cognee.tool")  # o ruta según cómo importes
    await tool.ingest(list(Path("data/acme").glob("*.md")))

    result = await tool.query(question)
    rprint("[bold cyan]Answer:[/bold cyan]", result["answer"])
    rprint(f"[dim]source={result['source']}  confidence={result['confidence']}[/dim]")

    if result["confidence"] < 0.5:
        lc.verify_fail()
        # un retry con la pregunta reformulada
        result2 = await tool.query(f"summarize what you know about: {question}")
        lc.retry()
        rprint("[bold cyan]Retry answer:[/bold cyan]", result2["answer"])
    else:
        lc.verify_ok()