Python based API example

This guide shows how to call AIRA from a workspace Jupyter notebook with the requests package, using the Breast Cancer Wisconsin (Diagnostic) dataset as example data.

AIRA provides an OpenAI-compatible chat completions API. A request is an HTTPS POST to <base URL>/chat/completions containing a model name and a list of messages, authenticated with the workspace API key as a bearer token.

This guide covers:

  1. Settings
  2. A helper for chat requests
  3. Connection check
  4. Load and summarise the dataset
  5. Ask AIRA to describe the dataset
  6. Ask AIRA to explain differences computed in pandas
  7. Request a JSON response
  8. Continue the conversation
  9. The thinking setting
Requirements
  • the WORKSPACE_API_KEY environment variable, set by the workspace
  • the Python packages requests and pandas
  • the dataset at /home/workspace/files/data/bcw.csv

The code expects the common CSV layout of this dataset: an id column, a diagnosis column with values M (malignant) and B (benign), and 30 numeric feature columns such as radius_mean. If your file uses a different outcome column, change LABEL_COL in the settings cell.

1. Settings

import json
import os
import re
import time

try:
    import pandas as pd
    import requests
except ImportError as exc:
    raise ImportError(
        f"Package '{exc.name}' is not installed in this kernel. "
        "Contact your workspace administrator."
    ) from exc

# AIRA endpoint. The default is the UK South endpoint. AIRA_BASE_URL is read
# by this notebook only, so the endpoint can be changed without editing code;
# the workspace does not set it.
BASE_URL = os.environ.get(
    "AIRA_BASE_URL", "https://api.uksouth.saas.aridhia.io/api/aira/v1"
).rstrip("/")
MODEL = "workspace-chat"

# Workspace API key, provided by the workspace. It is never printed.
API_KEY = os.environ.get("WORKSPACE_API_KEY", "")

DATA_PATH = "/home/workspace/files/data/bcw.csv"
LABEL_COL = "diagnosis"
LABEL_NAMES = {"M": "malignant", "B": "benign"}

TIMEOUT_S = 120  # seconds allowed per request

print("Endpoint:", BASE_URL)
print("Model:   ", MODEL)
print("API key: ", "set" if API_KEY else "NOT SET")

2. A helper for chat requests

Every call in this example goes through aira_chat. It sends the messages, measures the time taken and returns a dictionary instead of raising an exception, so a failed call does not stop the notebook. The status field is one of:

  • ok: the reply is in text
  • timeout: no response within timeout_s
  • unavailable: the connection failed or the endpoint returned an HTTP error
  • malformed: the response was not in the expected shape

The model served as workspace-chat has a thinking mode, in which it writes a reasoning block before its answer. disable_thinking=True (the default here) adds "chat_template_kwargs": {"enable_thinking": false} to the request body, which turns this off and shortens response times. If a reply does contain a <think> block, the helper removes it.

Other keyword arguments, such as temperature or max_tokens, are passed through in the request body.

THINK_BLOCK = re.compile(r"<think>.*?</think>", re.DOTALL)

def aira_chat(messages, disable_thinking=True, timeout_s=TIMEOUT_S, **params):
    """Send a chat completions request to AIRA and return a result dictionary."""
    body = {"model": MODEL, "messages": messages, **params}
    if disable_thinking:
        body["chat_template_kwargs"] = {"enable_thinking": False}

    result = {"status": "", "text": "", "reason": "", "duration_ms": 0,
              "finish_reason": None, "usage": None}
    start = time.monotonic()
    try:
        resp = requests.post(
            f"{BASE_URL}/chat/completions",
            headers={"Authorization": f"Bearer {API_KEY}"},
            json=body,
            timeout=timeout_s,
        )
    except requests.Timeout:
        result.update(status="timeout", reason=f"No response within {timeout_s} s.")
    except requests.RequestException as exc:
        result.update(status="unavailable", reason=f"{type(exc).__name__}: {exc}")
    else:
        if resp.status_code != 200:
            result.update(status="unavailable",
                          reason=f"HTTP {resp.status_code}: {resp.text[:300]}")
        else:
            try:
                data = resp.json()
                choice = data["choices"][0]
                content = choice["message"]["content"] or ""
            except (ValueError, KeyError, IndexError, TypeError) as exc:
                result.update(status="malformed",
                              reason=f"Unexpected response ({type(exc).__name__}).")
            else:
                result.update(status="ok",
                              text=THINK_BLOCK.sub("", content).strip(),
                              finish_reason=choice.get("finish_reason"),
                              usage=data.get("usage"))
    result["duration_ms"] = int((time.monotonic() - start) * 1000)
    return result

def show(result):
    """Print the timing line, then the reply or the reason for failure."""
    tokens = (result["usage"] or {}).get("completion_tokens")
    print(f"status={result['status']}  time={result['duration_ms'] / 1000:.1f} s  "
          f"finish_reason={result['finish_reason']}  completion_tokens={tokens}")
    print()
    print(result["text"] if result["status"] == "ok" else result["reason"])

3. Connection check

A short request confirms that the endpoint, model and key work before anything larger is sent.

if not API_KEY:
    print("WORKSPACE_API_KEY is not set, so the endpoint will refuse requests.\n")

check = aira_chat([{"role": "user", "content": "Reply with the single word: ready"}])
show(check)

4. Load and summarise the dataset

The data is read and summarised with pandas. Later sections send AIRA column names and summary statistics computed here, not individual rows.

df = pd.read_csv(DATA_PATH)

# Some copies of this file end with an empty unnamed column. Drop columns with no values.
df = df.dropna(axis=1, how="all")

has_label = LABEL_COL in df.columns
feature_cols = [c for c in df.select_dtypes("number").columns
                if c != LABEL_COL and c.lower() != "id"]

print(f"{len(df)} rows, {df.shape[1]} columns, {len(feature_cols)} numeric features")
print()
print(", ".join(df.columns))
print()
if has_label:
    counts = df[LABEL_COL].value_counts()
    class_counts = ", ".join(f"{k}: {int(v)}" for k, v in counts.items())
    print(f"{LABEL_COL}: {class_counts}")
else:
    class_counts = ""
    print(f"Column '{LABEL_COL}' not found. Set LABEL_COL in the settings cell.")

5. Ask AIRA to describe the dataset

The prompt contains the file's column names, row count and outcome counts, so the answer is based on the structure of the actual file. A system message sets the assistant's role for the conversation.

label_line = f"Counts of the {LABEL_COL} column: {class_counts}\n" if has_label else ""

describe_prompt = (
    f"A CSV file in a research workspace has {len(df)} rows and these columns:\n"
    f"{', '.join(df.columns)}\n"
    f"{label_line}\n"
    "Identify the dataset if you recognise it. Then describe in plain language what "
    "one row represents, what the outcome variable is, and how the feature columns "
    "are organised. Keep the answer under 200 words."
)

messages = [
    {"role": "system",
     "content": "You are a data analysis assistant in a research workspace. "
                "Answer factually and say when you are unsure."},
    {"role": "user", "content": describe_prompt},
]

described = aira_chat(messages)
show(described)

6. Ask AIRA to explain differences computed in pandas

Language models are not reliable at arithmetic over long tables. A more dependable pattern is to compute the figures in code and ask the model to explain them.

For each feature, pandas computes the mean in each group and Cohen's d: the difference between the group means divided by the pooled standard deviation. A positive d means the feature is higher in the malignant group. The ten features with the largest absolute d are sent to AIRA, which is asked to explain them using only the figures given.

effect = None
labels = sorted(df[LABEL_COL].dropna().unique()) if has_label else []

if len(labels) != 2:
    print("Skipped: this section needs an outcome column with exactly two values.")
else:
    # Order the groups so that d is (second group - first group). With M and B,
    # put benign first so a positive d means higher in the malignant group.
    g1, g2 = sorted(labels, key=lambda g: LABEL_NAMES.get(g, str(g)) == "malignant")
    name1, name2 = LABEL_NAMES.get(g1, str(g1)), LABEL_NAMES.get(g2, str(g2))

    grouped = df.groupby(LABEL_COL)[feature_cols]
    means, sds, ns = grouped.mean(), grouped.std(), grouped.size()
    pooled_sd = (((ns[g1] - 1) * sds.loc[g1] ** 2 + (ns[g2] - 1) * sds.loc[g2] ** 2)
                 / (ns[g1] + ns[g2] - 2)) ** 0.5

    effect = pd.DataFrame({
        f"mean_{name1}": means.loc[g1],
        f"mean_{name2}": means.loc[g2],
        "cohens_d": (means.loc[g2] - means.loc[g1]) / pooled_sd,
    })
    effect = effect.loc[effect["cohens_d"].abs().sort_values(ascending=False).index]
    display(effect.head(10).round(3))
if effect is not None:
    explain_prompt = (
        f"For a dataset of tumour measurements, these are the ten features that differ "
        f"most between the {name1} and {name2} groups ({int(ns[g1])} and {int(ns[g2])} rows). "
        f"cohens_d is the difference in means ({name2} minus {name1}) divided by the "
        f"pooled standard deviation.\n\n"
        f"{effect.head(10).round(4).to_csv()}\n"
        f"Using only these figures, explain in plain language how the two groups differ, "
        f"which kinds of measurement differ most, and one caution about interpreting "
        f"these results. Keep the answer under 250 words."
    )
    show(aira_chat([messages[0], {"role": "user", "content": explain_prompt}]))

7. Request a JSON response

For use in code, a reply can be requested as JSON with a fixed set of keys. Some models wrap JSON in Markdown code fences, so parse_json_reply removes them before parsing. The result is then checked: the required keys must be present, and every column the reply names must exist in the file.

temperature=0 is a standard OpenAI-compatible parameter, passed through by the helper, that makes the output less variable.

def parse_json_reply(text):
    """Remove Markdown code fences if present and parse the reply as JSON."""
    cleaned = text.strip()
    if cleaned.startswith("```"):
        cleaned = re.sub(r"^```[A-Za-z]*\s*|\s*```$", "", cleaned)
    return json.loads(cleaned)

REQUIRED_KEYS = {"dataset_name", "unit_of_observation", "outcome_column",
                 "feature_groups", "suggested_analyses"}

json_prompt = (
    f"Columns of a CSV file: {', '.join(df.columns)}\n\n"
    "Return only a JSON object, with no other text, in this shape:\n"
    "{\n"
    '  "dataset_name": "<name, or null if not recognised>",\n'
    '  "unit_of_observation": "<what one row represents>",\n'
    '  "outcome_column": "<column name, or null>",\n'
    '  "feature_groups": [{"name": "<group name>", "columns": ["<column name>"]}],\n'
    '  "suggested_analyses": ["<one sentence each, at most three>"]\n'
    "}"
)

structured = aira_chat([{"role": "user", "content": json_prompt}], temperature=0)
parsed = None

if structured["status"] != "ok":
    show(structured)
else:
    try:
        parsed = parse_json_reply(structured["text"])
    except json.JSONDecodeError as exc:
        print(f"The reply was not valid JSON ({exc}):\n")
        print(structured["text"])
    else:
        if not isinstance(parsed, dict):
            print("The reply was valid JSON but not an object.")
            parsed = None
        elif REQUIRED_KEYS - set(parsed):
            print("Missing keys:", sorted(REQUIRED_KEYS - set(parsed)))
            parsed = None

if parsed is not None:
    print(f"status=ok  time={structured['duration_ms'] / 1000:.1f} s\n")
    print("Dataset:            ", parsed["dataset_name"])
    print("Unit of observation:", parsed["unit_of_observation"])
    print("Outcome column:     ", parsed["outcome_column"])
    print("\nFeature groups:")
    groups = [g for g in parsed["feature_groups"] if isinstance(g, dict)]
    for g in groups:
        print(f"  {g.get('name')}: {len(g.get('columns', []))} columns")
    print("\nSuggested analyses:")
    for item in parsed["suggested_analyses"]:
        print(f"  - {item}")

    # Check the reply against the data.
    named = {c for g in groups for c in g.get("columns", [])}
    if parsed["outcome_column"]:
        named.add(parsed["outcome_column"])
    unknown = sorted(str(c) for c in named - set(df.columns))
    print("\nColumns named in the reply that are not in the file:", unknown or "none")

8. Continue the conversation

The API does not keep conversation state between requests. To ask a follow-up question, send the earlier messages again with the model's previous reply and the new question added to the end.

if described["status"] != "ok":
    print("Skipped: the request in section 5 did not succeed.")
else:
    follow_up = messages + [
        {"role": "assistant", "content": described["text"]},
        {"role": "user",
         "content": "Suggest a simple first model for predicting the outcome from these "
                    "features, and the main check needed to show it is not overfitting. "
                    "Keep the answer under 150 words."},
    ]
    show(aira_chat(follow_up))

9. The thinking setting

The same short question is sent twice, first with thinking disabled and then with it enabled, to compare response time and completion tokens. The second request is slower. Set RUN_THINKING_COMPARISON = False to skip this section.

RUN_THINKING_COMPARISON = True

if RUN_THINKING_COMPARISON:
    question = [{
        "role": "user",
        "content": "In one sentence, which of these measurements would you expect to "
                   "differ most between malignant and benign breast tumours: "
                   f"{', '.join(feature_cols[:10])}?",
    }]
    for disable in (True, False):
        print(f"disable_thinking={disable}")
        show(aira_chat(question, disable_thinking=disable))
        print("\n")

Notes

  • This example sends column names and summary statistics to AIRA, not individual rows. Whether row-level data may be included in prompts is a matter for the workspace's data governance arrangements.
  • Model replies can be wrong. Figures that matter should be computed in code, as in section 6, and structured output should be checked against the data, as in section 7.
  • Because the endpoint is OpenAI-compatible, other OpenAI-compatible clients can make the same requests by setting the base URL, the API key and the model name workspace-chat.
Updated on October 23, 2026
Back to top