We build intelligent software that solves real problems

# node59: intelligence, put to work
from __future__ import annotations
import json, math, time, hashlib
from dataclasses import dataclass, field
from pathlib import Path
import numpy as np
import torch, torch.nn as nn, torch.nn.functional as F

NODE = "node59"
EMBED_DIM = 768; CHUNK = 512; OVERLAP = 64; TOP_K = 8; SEED = 59
torch.manual_seed(SEED); np.random.seed(SEED)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def chunk_text(text: str, size: int = CHUNK, overlap: int = OVERLAP) -> list[str]:
    words = text.split(); step = max(1, size - overlap)
    return [" ".join(words[i:i + size]) for i in range(0, max(1, len(words)), step)]

def embed(texts: list[str], model) -> np.ndarray:
    with torch.no_grad(): vectors = model.encode(texts, batch_size=64, normalize_embeddings=True)
    return np.asarray(vectors, dtype=np.float32)

class VectorIndex:
    def __init__(self, dim: int = EMBED_DIM): self.dim, self.vectors, self.payloads = dim, np.empty((0, dim), np.float32), []
    def add(self, vectors: np.ndarray, payloads: list[dict]) -> None:
        self.vectors = np.vstack([self.vectors, vectors]); self.payloads.extend(payloads)
    def search(self, query: np.ndarray, k: int = TOP_K) -> list[tuple[float, dict]]:
        scores = self.vectors @ query; order = np.argsort(-scores)[:k]
        return [(float(scores[i]), self.payloads[i]) for i in order]
    def save(self, path: Path) -> None: np.save(path / f"{NODE}.npy", self.vectors); (path / f"{NODE}.json").write_text(json.dumps(self.payloads))

class Classifier(nn.Module):
    def __init__(self, dim: int = EMBED_DIM, hidden: int = 256, classes: int = 12, dropout: float = 0.1):
        super().__init__(); self.net = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, classes))
    def forward(self, x: torch.Tensor) -> torch.Tensor: return self.net(x)

def train(model: nn.Module, loader, epochs: int = 12, lr: float = 3e-4) -> list[float]:
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01); history = []
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs * len(loader))
    for epoch in range(epochs):
        model.train(); total = 0.0
        for x, y in loader:
            x, y = x.to(device), y.to(device); optimizer.zero_grad(set_to_none=True)
            loss = F.cross_entropy(model(x), y, label_smoothing=0.05); loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 1.0); optimizer.step(); scheduler.step(); total += loss.item()
        history.append(total / len(loader)); print(f"[{NODE}] epoch {epoch + 1:02d} loss {history[-1]:.4f}")
    return history

@torch.no_grad()
def evaluate(model: nn.Module, loader) -> dict[str, float]:
    model.eval(); correct = seen = 0
    for x, y in loader: correct += (model(x.to(device)).argmax(-1) == y.to(device)).sum().item(); seen += len(y)
    return {"accuracy": correct / max(1, seen), "seen": seen}

def rerank(query: str, hits: list[tuple[float, dict]], scorer) -> list[dict]:
    pairs = [(query, hit["text"]) for _, hit in hits]; scores = scorer.predict(pairs)
    return [hit for _, hit in sorted(zip(scores, [h for _, h in hits]), key=lambda s: -s[0])]

def build_prompt(question: str, context: list[dict]) -> str:
    cited = "\n\n".join(f"[{i + 1}] {c['source']}: {c['text']}" for i, c in enumerate(context))
    return f"Answer using only the context. Cite sources as [n].\n\nContext:\n{cited}\n\nQuestion: {question}\nAnswer:"

@dataclass
class Node:
    name: str; run: callable; inputs: list[str] = field(default_factory=list); retries: int = 2
    def __call__(self, state: dict) -> dict:
        for attempt in range(self.retries + 1):
            try: return {**state, **self.run(state)}
            except Exception as error:
                if attempt == self.retries: raise
                time.sleep(0.25 * (attempt + 1))

class NodeGraph:
    def __init__(self, name: str = NODE): self.name, self.nodes, self.edges = name, {}, {}
    def add(self, node: Node, after: list[str] | None = None) -> "NodeGraph":
        self.nodes[node.name] = node; self.edges[node.name] = list(after or []); return self
    def order(self) -> list[str]:
        seen, out = set(), []
        def visit(n): [visit(d) for d in self.edges[n] if d not in seen]; seen.add(n); out.append(n)
        for n in self.nodes:
            if n not in seen: visit(n)
        return out
    def run(self, state: dict) -> dict:
        for name in self.order(): state = self.nodes[name](state); state["trace"] = state.get("trace", []) + [name]
        return state

def fingerprint(record: dict) -> str: return hashlib.sha256(json.dumps(record, sort_keys=True).encode()).hexdigest()[:16]

def retrieve(state: dict) -> dict:
    query = embed([state["question"]], state["encoder"])[0]; hits = state["index"].search(query, k=TOP_K)
    return {"hits": hits, "context": rerank(state["question"], hits, state["scorer"])[:4]}

def answer(state: dict) -> dict:
    prompt = build_prompt(state["question"], state["context"]); reply = state["llm"].complete(prompt, temperature=0.1, max_tokens=600)
    return {"answer": reply.text, "citations": [c["source"] for c in state["context"]], "fingerprint": fingerprint({"q": state["question"], "a": reply.text})}

def classify(state: dict) -> dict:
    logits = state["classifier"](torch.from_numpy(embed([state["question"]], state["encoder"])).to(device))
    return {"intent": state["labels"][int(logits.argmax(-1))], "confidence": float(F.softmax(logits, -1).max())}

graph = NodeGraph(NODE).add(Node("classify", classify)).add(Node("retrieve", retrieve), after=["classify"]).add(Node("answer", answer), after=["retrieve"])

if __name__ == "__main__":
    index = VectorIndex(); index.add(embed(chunks, encoder), [{"source": src, "text": chunk} for src, chunk in zip(sources, chunks)])
    history = train(Classifier().to(device), train_loader); print(f"[{NODE}] {evaluate(classifier, eval_loader)}")
    result = graph.run({"question": "Which invoices are overdue by more than 30 days?", "index": index, "encoder": encoder, "scorer": scorer, "llm": llm, "classifier": classifier, "labels": LABELS})
    print(json.dumps({k: result[k] for k in ("intent", "confidence", "answer", "citations", "trace")}, indent=2))