Managed LLM & RAG

import os, shutil, subprocess

print("OPENAI_API_BASE =", os.environ.get("OPENAI_API_BASE", "NOT SET"))
# Print the actual token. This is a shared workshop key — fine to show here;
# treat your own personal token as a secret.
key = os.environ.get("OPENAI_API_KEY", "")
print("OPENAI_API_KEY  =", key if key else "NOT SET")

kubectl = shutil.which("kubectl") or "/opt/conda/bin/kubectl"
print("kubectl path    =", kubectl)
print("kubectl version =", subprocess.run([kubectl, "version", "--client", "-o", "json"],
                                          capture_output=True, text=True).stdout[:120], "...")

print("can list pods in nrp-training-k8s:",
      subprocess.run([kubectl, "auth", "can-i", "list", "pods", "-n", "nrp-training-k8s"],
                     capture_output=True, text=True).stdout.strip())

# List the models currently served by the NRP managed endpoint.
import os, json
from openai import OpenAI

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"], base_url=os.environ["OPENAI_API_BASE"])
models = client.models.list()
print(f"{len(models.data)} models available right now:\n")
for m in sorted(models.data, key=lambda x: x.id):
    print(f"  {m.id}")

# A tiny helper we'll reuse across this notebook. The default model is
# `gemma-small-e4b` (Gemma 3n E4B — fast on NRP and multimodal). The helper also
# handles *reasoning* models like `minimax-m2` (you'll try one below): they
# stream a private chain-of-thought into a separate `reasoning` field and only
# fill `content` once they conclude — so we default `max_tokens` high enough
# (1200) to leave room for both the thinking and the final answer. If a call
# still runs out of tokens mid-thought, we surface that reasoning rather than
# returning nothing.
def chat(prompt, model="gemma-small-e4b", system=None, max_tokens=1200, llm=None):
    msgs = []
    if system:
        msgs.append({"role": "system", "content": system})
    msgs.append({"role": "user", "content": prompt})
    resp = (llm or client).chat.completions.create(
        model=model, messages=msgs, max_tokens=max_tokens, temperature=0.2,
    )
    msg = resp.choices[0].message
    if msg.content:
        return msg.content
    reasoning = getattr(msg, "reasoning", None) or getattr(msg, "reasoning_content", None)
    if reasoning:
        return f"(no final answer within max_tokens={max_tokens}; reasoning so far:)\n{reasoning}"
    return "(model produced no content — try raising max_tokens or another model)"

# Single chat completion.
print(chat(
    "What is the National Research Platform?",
    system="Answer in one sentence.",
))

# The system prompt sets the model's ROLE. Same question, four useful roles:
QUESTION = ("How do I compute the average and standard deviation of a list of "
            "student grades in Python?")

ROLES = {
    "Teaching assistant": (
        "You are a supportive teaching assistant. Explain clearly and build "
        "intuition with a short example, guiding the learner toward the answer "
        "rather than just dumping it."),
    "Technical coder": (
        "You are an expert software engineer. Write clean, correct, "
        "well-commented Python, then briefly explain it and note any edge cases."),
    "Concise expert": (
        "You are a domain expert. Give a precise, rigorous answer in a few "
        "sentences for a graduate-level audience. No filler."),
    "Research assistant": (
        "You are a meticulous research assistant. Give a structured answer, "
        "separate established facts from uncertainty, and point to where to "
        "learn more."),
}
for role, system in ROLES.items():
    print(f"=== {role} ===")
    print(chat(QUESTION, system=system, model="gemma-small-e4b"), "\n")

# A simple interactive chat — works in ANY session (no widget extension needed).
# Set ROLE + MODEL, run the cell, then chat. Type 'quit' to stop, 'reset' to clear.
SYSTEMS = {
    "Teaching assistant": "You are a supportive teaching assistant. Explain clearly, "
        "build intuition with short examples, and guide the learner toward the answer.",
    "Technical coder": "You are an expert software engineer. Write clean, correct, "
        "well-commented code and briefly explain it.",
    "Concise expert": "You are a domain expert. Answer precisely in a few sentences, "
        "graduate level, no filler.",
    "Research assistant": "You are a meticulous research assistant. Give structured "
        "answers, flag uncertainty, and suggest where to learn more.",
}
ROLE  = "Teaching assistant"   # <- change to any key above
MODEL = "gemma-small-e4b"          # <- or minimax-m2, gpt-oss, qwen3, gemma

history = []
print(f"Chatting as: {ROLE} ({MODEL}).  Type 'quit' to stop, 'reset' to clear.\n")
while True:
    try:
        msg = input("You: ").strip()
    except (EOFError, KeyboardInterrupt):
        print("\n(ended)"); break
    if msg.lower() in ("quit", "exit", "q", ""):
        print("Bye!"); break
    if msg.lower() == "reset":
        history.clear(); print("(conversation cleared)\n"); continue
    history.append({"role": "user", "content": msg})
    r = client.chat.completions.create(
        model=MODEL, max_tokens=1000, temperature=0.5,
        messages=[{"role": "system", "content": SYSTEMS[ROLE]}] + history)
    m = r.choices[0].message
    reply = m.content or getattr(m, "reasoning", None) or "(no reply)"
    history.append({"role": "assistant", "content": reply})
    print(f"AI: {reply}\n")

import numpy as np, matplotlib.pyplot as plt

def embed(texts):
    r = client.embeddings.create(model="qwen3-embedding", input=texts)
    v = np.array([d.embedding for d in r.data])
    return v / np.linalg.norm(v, axis=1, keepdims=True)

docs = [
    "A GPU accelerates deep-learning training.",
    "Kubernetes schedules containers across a cluster.",
    "Helm is a package manager for Kubernetes.",
    "The recipe needs two cups of flour.",
    "Cats like to nap in the sun.",
]
D = embed(docs)
query = "How do I run AI workloads on a compute cluster?"
sims = D @ embed([query])[0]
print(f"Query: {query}\n")
for i in sims.argsort()[::-1]:
    print(f"  {sims[i]:.3f}  {docs[i]}")

M = D @ D.T
fig, ax = plt.subplots(figsize=(5, 4))
im = ax.imshow(M, cmap="viridis", vmin=0, vmax=1)
lab = [d[:20] + "..." for d in docs]
ax.set_xticks(range(len(docs))); ax.set_xticklabels(lab, rotation=45, ha="right", fontsize=7)
ax.set_yticks(range(len(docs))); ax.set_yticklabels(lab, fontsize=7)
fig.colorbar(im, label="cosine similarity")
ax.set_title("Sentence similarity (qwen3-embedding)")
plt.tight_layout(); plt.show()

import base64, requests
from IPython.display import Image, display

IMG_URL = "https://picsum.photos/id/237/420/280"     # any image URL works
raw = requests.get(IMG_URL, timeout=20).content
display(Image(data=raw))                              # show it inline

b64 = base64.b64encode(raw).decode()
r = client.chat.completions.create(model="gemma-small-e4b", max_tokens=200, messages=[
    {"role": "user", "content": [
        {"type": "text", "text": "Describe this image — what's in it, and the mood?"},
        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
    ]}])
print("gemma-small-e4b sees:\n", r.choices[0].message.content)

# Pull one real NRP documentation page (the cluster policies page) and clean it.
import requests, re

PAGE = "https://nrp.ai/documentation/userdocs/start/policies/"
RAW  = ("https://gitlab.nrp-nautilus.io/prp/nrp-site/-/raw/main/"
        "src/content/docs/Documentation/userdocs/start/policies.mdx")

md = requests.get(RAW, timeout=30).text
md = re.sub(r"^---.*?---\s*", "", md, flags=re.S)        # drop frontmatter
md = re.sub(r"^import .*$", "", md, flags=re.M)          # drop MDX imports
md = re.sub(r":::\w+(\[[^\]]*\])?|:::", "", md)          # drop admonition markers
md = re.sub(r"<[^>]+>", "", md)                          # drop stray HTML/JSX
md = re.sub(r"\n{3,}", "\n\n", md).strip()

# Split into ~700-char chunks. (A real corpus would chunk every page; one is plenty here.)
chunks = [md[i:i+700] for i in range(0, len(md), 550)]
print(f"Pulled {len(md)} chars from {PAGE}")
print(f"Split into {len(chunks)} chunks.")

# Embed with NRP's MANAGED embedding model — no local model, nothing to download.
# Same `client`, same endpoint as the chat models; just a different call.
import numpy as np

def embed(texts):
    resp = client.embeddings.create(model="qwen3-embedding", input=texts)
    vecs = np.array([d.embedding for d in resp.data])
    return vecs / np.linalg.norm(vecs, axis=1, keepdims=True)   # normalize for cosine

chunk_vecs = embed(chunks)          # one API call embeds the whole page
print(f"Embedded {len(chunks)} chunks into vectors of dim {chunk_vecs.shape[1]} (via qwen3-embedding).")

def retrieve(question, k=3):
    qv = embed([question])[0]
    scores = chunk_vecs @ qv         # cosine similarity (everything is normalized)
    top = scores.argsort()[-k:][::-1]
    return [(chunks[i], float(scores[i])) for i in top]

# Answer a question by RAG, using NRP's managed LLM.
SYSTEM = ("Answer the question using ONLY the provided context. "
          "If the context doesn't contain the answer, say so. Be concise.")

def ask(question, model="minimax-m2"):
    context = "\n\n".join(text for text, _ in retrieve(question))
    return chat(f"Context:\n{context}\n\nQuestion: {question}", system=SYSTEM, model=model)

QUESTION = "Is it okay to run `sleep infinity` in a Job on NRP?"
print("Retrieved chunks:")
for text, score in retrieve(QUESTION):
    print(f"  score={score:.3f}  {text[:70].strip()}...")

print("\n--- Answer (NRP managed minimax-m2) ---")
print(ask(QUESTION))