cortex_client.py

"""
Cortex client mirroring ATLAS patterns.
Uses HTTP REST API (POST /nodes, POST /edges, GET /nodes, POST /search, etc.)
Based on: https://github.com/MikeSquared-Agency/ATLAS/blob/main/atlas/cortex.py
"""

import httpx
import json
from dataclasses import dataclass, field
from typing import Any


# ─── Data Classes (matching ATLAS) ───

@dataclass
class CortexNode:
    """Represents a node in Cortex."""
    id: str | None  # None for new nodes, set after creation
    kind: str
    title: str
    body: dict[str, Any]  # stored as JSON string in Cortex
    tags: list[str] = field(default_factory=list)
    source_agent: str = "recursive-ml-researcher"
    importance: float = 0.5


@dataclass
class CortexEdge:
    """Represents an edge in Cortex."""
    from_id: str
    to_id: str
    relation: str
    weight: float = 1.0


# ─── Helper Functions ───

def _check_response(data: dict) -> None:
    """Raise on Cortex error responses."""
    if isinstance(data, dict) and data.get("error"):
        raise RuntimeError(f"Cortex API error: {data['error']}")


def _node_from_response(data: dict) -> CortexNode:
    """Parse a CortexNode from API response data."""
    body = data.get("body", "{}")
    if isinstance(body, str):
        try:
            body = json.loads(body)
        except (json.JSONDecodeError, TypeError):
            body = {"raw": body}
    elif not isinstance(body, dict):
        body = {"raw": str(body)}

    return CortexNode(
        id=data.get("id"),
        kind=data.get("kind", "unknown"),
        title=data.get("title", ""),
        body=body,
        tags=data.get("tags", []),
        source_agent=data.get("source_agent", ""),
        importance=data.get("importance", 0.5),
    )


# ─── Cortex Client ───

class CortexClient:
    """Async HTTP client for Cortex API.

    Mirrors the ATLAS pattern:
      - POST /nodes          create node
      - PATCH /nodes/:id     update node
      - POST /edges          create edge
      - GET  /nodes          list nodes (params: kind, tag, limit)
      - GET  /nodes/:id      get single node
      - GET  /nodes/:id/neighbors  get adjacent nodes
      - POST /search         semantic search
      - GET  /graph/export   export full graph
      - POST /auto-linker/trigger  trigger auto-linking
    """

    def __init__(self, base_url: str = "http://localhost:9091") -> None:
        self.base_url = base_url.rstrip("/")
        self.session_scope = None
        self._warned = False
        self._node_ids: dict[str, str] = {}  # title -> id cache

        try:
            self._client = httpx.AsyncClient(
                base_url=self.base_url,
                timeout=30.0,
                headers={"x-gate-override": "true"},
            )
        except Exception as e:
            print(f"WARNING: Could not create Cortex HTTP client: {e}")
            self._client = None

    # ─── Node Operations ───

    async def create_node(self, node: CortexNode) -> CortexNode:
        """Create a node. Returns the node with id set."""
        if not self._client:
            return self._fallback_node(node)

        payload = {
            "kind": node.kind,
            "title": node.title,
            "body": json.dumps(node.body),
            "tags": node.tags,
            "source_agent": node.source_agent,
            "importance": node.importance,
        }
        try:
            resp = await self._client.post("/nodes", params={"gate": "skip"}, json=payload)
            data = resp.json()
            _check_response(data)
            result = _node_from_response(data.get("data", data))
            if result.id:
                self._node_ids[node.title] = result.id
            return result
        except Exception as e:
            self._warn("create_node", e)
            return self._fallback_node(node)

    async def update_node(self, node_id: str, **fields: Any) -> CortexNode | None:
        """Update node fields. Only pass fields that changed."""
        if not self._client:
            return None

        payload: dict[str, Any] = {}
        for key, value in fields.items():
            if key == "body" and isinstance(value, dict):
                payload["body"] = json.dumps(value)
            else:
                payload[key] = value

        try:
            resp = await self._client.patch(f"/nodes/{node_id}", params={"gate": "skip"}, json=payload)
            data = resp.json()
            _check_response(data)
            return _node_from_response(data.get("data", data))
        except Exception as e:
            self._warn("update_node", e)
            return None

    async def find_node_by_title(self, title: str) -> CortexNode | None:
        """Find an existing node by its title."""
        if not self._client:
            return None
        try:
            resp = await self._client.get("/nodes", params={"title": title, "limit": 1})
            data = resp.json()
            nodes = data.get("data", data.get("nodes", []))
            if nodes and len(nodes) > 0:
                return _node_from_response(nodes[0])
        except Exception:
            pass
        return None

    async def upsert_node(self, node: CortexNode) -> CortexNode:
        """Create or update a node. Matches by title (ATLAS pattern)."""
        existing = await self.find_node_by_title(node.title)
        if existing and existing.id:
            updated = await self.update_node(
                existing.id,
                body=node.body,
                tags=node.tags,
                importance=node.importance,
            )
            return updated or existing
        return await self.create_node(node)

    # ─── Edge Operations ───

    async def create_edge(self, edge: CortexEdge) -> None:
        """Create an edge between two nodes."""
        if not self._client:
            return

        payload = {
            "from_id": edge.from_id,
            "to_id": edge.to_id,
            "relation": edge.relation,
            "weight": edge.weight,
        }
        try:
            resp = await self._client.post("/edges", params={"gate": "skip"}, json=payload)
            data = resp.json()
            _check_response(data)
        except Exception as e:
            self._warn("create_edge", e)

    # ─── Query Operations ───

    async def list_nodes(self, kind: str = None, tag: str = None, limit: int = 100) -> list[CortexNode]:
        """List nodes, optionally filtered by kind and/or tag."""
        if not self._client:
            return []

        params: dict[str, Any] = {"limit": limit}
        if kind:
            params["kind"] = kind
        if tag:
            params["tag"] = tag

        try:
            resp = await self._client.get("/nodes", params=params)
            data = resp.json()
            nodes = data.get("data", data.get("nodes", []))
            return [_node_from_response(n) for n in nodes]
        except Exception as e:
            self._warn("list_nodes", e)
            return []

    async def get_node(self, node_id: str) -> CortexNode | None:
        """Get a single node by ID."""
        if not self._client:
            return None
        try:
            resp = await self._client.get(f"/nodes/{node_id}")
            data = resp.json()
            return _node_from_response(data.get("data", data))
        except Exception as e:
            self._warn("get_node", e)
            return None

    async def get_neighbors(self, node_id: str, depth: int = 1) -> list[CortexNode]:
        """Get adjacent nodes."""
        if not self._client:
            return []
        try:
            resp = await self._client.get(f"/nodes/{node_id}/neighbors", params={"depth": depth})
            data = resp.json()
            nodes = data.get("data", [])
            return [_node_from_response(n) for n in nodes]
        except Exception as e:
            self._warn("get_neighbors", e)
            return []

    async def search(self, query: str, limit: int = 10) -> list[CortexNode]:
        """Semantic search for related nodes."""
        if not self._client:
            return []
        try:
            resp = await self._client.post("/search", json={"query": query, "limit": limit})
            data = resp.json()
            nodes = data.get("data", data.get("results", []))
            return [_node_from_response(n) for n in nodes]
        except Exception as e:
            self._warn("search", e)
            return []

    # ─── Graph Export ───

    async def export_graph(self) -> dict:
        """Export full graph as JSON. Returns {"nodes": [...], "edges": [...]}."""
        if not self._client:
            return self._fallback_graph()
        try:
            resp = await self._client.get("/graph/export")
            if resp.status_code == 200:
                data = resp.json()
                return data.get("data", data)
        except Exception as e:
            self._warn("export_graph", e)
        return self._fallback_graph()

    async def trigger_auto_linker(self) -> None:
        """Trigger Cortex auto-linking to discover relationships."""
        if not self._client:
            return
        try:
            await self._client.post("/auto-linker/trigger")
        except Exception as e:
            self._warn("trigger_auto_linker", e)

    async def get_graph_stats(self) -> dict:
        """Get node/edge counts by type for UI display."""
        graph = await self.export_graph()
        nodes = graph.get("nodes", [])
        edges = graph.get("edges", [])

        node_counts = {}
        for n in nodes:
            kind = n.get("kind", "unknown") if isinstance(n, dict) else getattr(n, "kind", "unknown")
            node_counts[kind] = node_counts.get(kind, 0) + 1

        return {
            "total_nodes": len(nodes),
            "total_edges": len(edges),
            "by_kind": node_counts
        }

    # ─── Convenience: High-Level Store (backward-compatible with research_loop) ───

    async def store(self, kind: str, title: str, body: str = "", metadata: dict = None, importance: float = 0.7) -> str:
        """High-level store that creates a CortexNode. Returns node ID."""
        node_body = metadata.copy() if metadata else {}
        if body:
            node_body["description"] = body

        tags = [kind]
        if self.session_scope:
            tags.append(f"session-{self.session_scope}")

        node = CortexNode(
            id=None,
            kind=kind,
            title=title,
            body=node_body,
            tags=tags,
            importance=importance,
        )
        result = await self.upsert_node(node)
        return result.id or f"fallback-{kind}-{title[:20]}"

    async def store_and_link(self, kind: str, title: str, body: str = "",
                             metadata: dict = None, importance: float = 0.7,
                             link_from: str = None, relation: str = "related_to") -> str:
        """Store a node and optionally create an edge from another node."""
        node_id = await self.store(kind, title, body, metadata, importance)

        if link_from and node_id and not node_id.startswith("fallback-"):
            await self.create_edge(CortexEdge(
                from_id=link_from,
                to_id=node_id,
                relation=relation,
            ))

        return node_id

    # ─── Internal Helpers ───

    def _fallback_node(self, node: CortexNode) -> CortexNode:
        """Return a node with a fallback ID when Cortex is unavailable."""
        fallback_id = f"fallback-{node.kind}-{len(self._node_ids)}"
        node.id = fallback_id
        self._node_ids[node.title] = fallback_id
        return node

    def _fallback_graph(self) -> dict:
        """Build a graph from cached node IDs when Cortex is unavailable."""
        nodes = [{"id": nid, "title": title} for title, nid in self._node_ids.items()]
        return {"nodes": nodes, "edges": []}

    def _warn(self, method: str, error: Exception) -> None:
        """Print a warning without crashing."""
        if not self._warned:
            print(f"WARNING: Cortex not available ({method}): {error}")
            print("  The app will continue without Cortex. Start Cortex with:")
            print("  docker run -d -p 9090:9090 -p 9091:9091 mikesquared/cortex:latest")
            self._warned = True