per-customer cost is unknowable because tenants share a key

"""Report whether OpenAI usage can be attributed to your customers at all.

Read only. GET requests and nothing else: OPENAI_ADMIN_KEY must be an
organization admin key (sk-admin-...) with read scopes.

The finding here is not a number, it is a fact about the reporting dimensions.
user_id on the Usage API is the org member or service account that owns the
calling API key. It is never an end-user identifier you supplied, and the
request-level `user` field does not reach this endpoint at all. So the repair is
architectural, it is forward-only, and it is printed rather than performed:
minting credentials for your tenants is not an audit's job.
"""
import argparse
import datetime as dt
import logging
import os
import sys

import requests

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("openai_tenant_attribution_audit")

API = "https://api.openai.com/v1"

# The complete list of dimensions the platform can attribute along. Not a
# starting point: there is no fourth one, and that is the note.
DIMENSIONS = ("user_id", "api_key_id", "project_id")

FINDINGS = ("single-key", "keys-below-tenants")


def fold(pages):
    """Sum usage into the three dimensions the platform actually holds. Pure.

    Returns {"users": {id: requests}, "keys": {...}, "projects": {...},
    "requests": total}. Buckets with a null grouping value are counted into the
    total but not into a dimension, because a null there means "not attributed"
    and inventing a bucket for it would flatter the result.
    """
    out = {"users": {}, "keys": {}, "projects": {}, "requests": 0}
    for page in pages:
        for bucket in page.get("data") or []:
            for result in bucket.get("results") or []:
                try:
                    n = int(result.get("num_model_requests") or 0)
                except (TypeError, ValueError):
                    n = 0
                out["requests"] += n
                for field, key in (("user_id", "users"), ("api_key_id", "keys"),
                                   ("project_id", "projects")):
                    value = result.get(field)
                    if value:
                        name = str(value)
                        out[key][name] = out[key].get(name, 0) + n
    return out


def classify(user_id, directory):
    """What kind of principal is this user_id? Pure.

    "service-account", "member", or "unresolved". The first two are the same
    finding wearing different clothes: both are your own principals, neither is
    a customer. The third is a different problem entirely.
    """
    entry = directory.get(str(user_id))
    if entry is None:
        return "unresolved"
    if entry.get("service_account"):
        return "service-account"
    return "member"


def unresolved(folded, directory):
    """user_ids generating usage that the org directory does not know. Pure.

    Sorted, so two runs print the same order. Usually empty; when it is not,
    something is calling the API as a principal nobody can name, which wants
    answering before the attribution question does.
    """
    return sorted(u for u in folded.get("users", {})
                  if classify(u, directory) == "unresolved")


def verdict(folded, directory, tenant_count=None):
    """Can this organization's usage be sliced per customer? Pure.

    Returns (state, detail). tenant_count comes from your database because the
    API has no concept of a tenant; without it the script can still report the
    cardinality and the fact that every principal is one of your own, which is
    most of the answer.
    """
    keys = folded.get("keys") or {}
    users = folded.get("users") or {}
    total = folded.get("requests") or 0

    if total <= 0 and not keys:
        return ("no-usage",
                "no completions usage in the window, so there is nothing to "
                "attribute yet")

    kinds = sorted({classify(u, directory) for u in users})
    principal_note = ("%d user_id value(s), all of them org members or service "
                      "accounts rather than customers" % len(users))
    if "unresolved" in kinds:
        principal_note = ("%d user_id value(s), of which some resolve to nobody "
                          "in the org directory" % len(users))

    if len(keys) == 1:
        return ("single-key",
                "1 api_key_id covers every request in the window. There is one "
                "bucket, so per-customer cost has no place to come from. %s."
                % principal_note)

    if tenant_count is None:
        return ("unknown-tenant-count",
                "%d distinct api_key_id value(s) and %d project(s). %s. Pass "
                "the tenant count to judge whether that is enough buckets."
                % (len(keys), len(folded.get("projects") or {}), principal_note))

    if len(keys) < tenant_count:
        return ("keys-below-tenants",
                "%d distinct api_key_id value(s) against %d tenant(s). Cost per "
                "customer is unrecoverable by construction: the finest slice the "
                "platform can offer is one key, and there are fewer keys than "
                "customers. %s." % (len(keys), tenant_count, principal_note))

    return ("segmented",
            "%d distinct api_key_id value(s) for %d tenant(s), so the platform "
            "can slice finely enough. Confirm your key-to-tenant map is current."
            % (len(keys), tenant_count))


def get(session, path, params=None):
    r = session.get(API + path, params=params or {}, timeout=60)
    if r.status_code == 401:
        raise SystemExit("401 from OpenAI: OPENAI_ADMIN_KEY must be an "
                         "organization admin key, not a project key")
    if r.status_code == 403:
        raise SystemExit("403 from OpenAI: the key is not authorised for "
                         "/v1/organization")
    r.raise_for_status()
    return r.json()


def usage_pages(session, start_time, days, max_pages=20):
    params = {"start_time": start_time, "bucket_width": "1d", "limit": days,
              "group_by": list(DIMENSIONS)}
    for _ in range(max_pages):
        page = get(session, "/organization/usage/completions", params)
        yield page
        cursor = page.get("next_page")
        if not cursor:
            return
        params = dict(params, page=cursor)


def org_directory(session, max_pages=20):
    """The org's own principals, keyed by id, from GET /v1/organization/users."""
    out = {}
    params = {"limit": 100}
    for _ in range(max_pages):
        page = get(session, "/organization/users", params)
        data = page.get("data") or []
        for user in data:
            out[str(user.get("id"))] = {
                "name": user.get("name") or user.get("email") or "?",
                "service_account": bool(user.get("is_service_account")),
            }
        if not page.get("has_more") or not data:
            break
        params = {"limit": 100, "after": data[-1].get("id")}
    return out


def spend_by_key(session, start_time):
    out = {}
    page = get(session, "/organization/costs",
               {"start_time": start_time, "limit": 30, "group_by": "api_key_id"})
    for bucket in page.get("data") or []:
        for result in bucket.get("results") or []:
            key_id = str(result.get("api_key_id") or "unattributed")
            amount = (result.get("amount") or {}).get("value") or 0
            try:
                out[key_id] = out.get(key_id, 0.0) + float(amount)
            except (TypeError, ValueError):
                pass
    return out


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--tenants", type=int, default=None,
                    help="how many customers you serve; comes from your database")
    ap.add_argument("--days", type=int, default=7,
                    help="days of usage to fold (default 7)")
    args = ap.parse_args()

    key = os.environ.get("OPENAI_ADMIN_KEY")
    if not key:
        log.error("set OPENAI_ADMIN_KEY (an organization admin key with read scopes)")
        return 2

    session = requests.Session()
    session.headers.update({"Authorization": "Bearer " + key})

    now = dt.datetime.now(dt.timezone.utc)
    start = int((now - dt.timedelta(days=args.days)).timestamp())

    folded = fold(usage_pages(session, start, args.days))
    directory = org_directory(session)
    state, detail = verdict(folded, directory, args.tenants)

    log.info("%-20s %s", state, detail)

    for user_id in sorted(folded["users"], key=lambda u: -folded["users"][u]):
        kind = classify(user_id, directory)
        name = directory.get(user_id, {}).get("name", "not in the directory")
        log.info("  principal %-30s %-16s %s", user_id, kind, name)

    orphans = unresolved(folded, directory)
    if orphans:
        log.warning("  %d user_id value(s) resolve to nobody in the org "
                    "directory: %s", len(orphans), ", ".join(orphans))
        log.warning("  repair: find what is calling as these principals before "
                    "you touch the attribution question")

    if state in FINDINGS:
        spend = spend_by_key(session, int((now - dt.timedelta(days=30)).timestamp()))
        for key_id, amount in sorted(spend.items(), key=lambda kv: -kv[1])[:10]:
            log.warning("  30d spend  %-30s $%.2f", key_id, amount)
        log.warning("  repair: the Usage API cannot segment by end user. Mint "
                    "one key, or one project, per tenant or tenant tier via "
                    "/v1/organization/projects/{id}/service_accounts/{id}/api_keys "
                    "and attribute with group_by=api_key_id.")
        log.warning("  repair: this is forward-only and cannot backfill. Until "
                    "then, record each response's usage block against your own "
                    "tenant id and reconcile it against /v1/organization/costs.")
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())

/**
 * Report whether OpenAI usage can be attributed to your customers at all.
 *
 * Read only. GET requests and nothing else: OPENAI_ADMIN_KEY must be an
 * organization admin key with read scopes.
 *
 * The finding is a fact about the reporting dimensions, not a number. user_id
 * on the Usage API is the org member or service account that owns the calling
 * key, never an end-user identifier you supplied. The repair is architectural,
 * forward-only, and printed rather than performed.
 */
const API = 'https://api.openai.com/v1';

// The complete list of dimensions the platform can attribute along.
const DIMENSIONS = ['user_id', 'api_key_id', 'project_id'];

const FINDINGS = ['single-key', 'keys-below-tenants'];

/**
 * Sum usage into the three dimensions the platform actually holds. Pure.
 * A null grouping value counts into the total but into no dimension, because
 * null means "not attributed" and inventing a bucket would flatter the result.
 */
export function fold(pages) {
  const out = { users: {}, keys: {}, projects: {}, requests: 0 };
  for (const page of pages) {
    for (const bucket of page.data ?? []) {
      for (const result of bucket.results ?? []) {
        const raw = Number(result.num_model_requests ?? 0);
        const n = Number.isFinite(raw) ? Math.trunc(raw) : 0;
        out.requests += n;
        for (const [field, key] of [['user_id', 'users'], ['api_key_id', 'keys'],
                                    ['project_id', 'projects']]) {
          const value = result[field];
          if (value) {
            const name = String(value);
            out[key][name] = (out[key][name] ?? 0) + n;
          }
        }
      }
    }
  }
  return out;
}

/**
 * What kind of principal is this user_id? Pure. "service-account", "member" or
 * "unresolved". The first two are the same finding in different clothes.
 */
export function classify(userId, directory) {
  const entry = directory[String(userId)];
  if (entry === undefined || entry === null) return 'unresolved';
  return entry.service_account ? 'service-account' : 'member';
}

/** user_ids generating usage that the org directory does not know. Pure. */
export function unresolved(folded, directory) {
  return Object.keys(folded.users ?? {})
    .filter((u) => classify(u, directory) === 'unresolved')
    .sort();
}

/**
 * Can this organization's usage be sliced per customer? Pure.
 * Returns [state, detail]. tenantCount comes from your database, because the
 * API has no concept of a tenant.
 */
export function verdict(folded, directory, tenantCount = null) {
  const keys = folded.keys ?? {};
  const users = folded.users ?? {};
  const total = folded.requests ?? 0;
  const keyCount = Object.keys(keys).length;
  const userCount = Object.keys(users).length;

  if (total <= 0 && keyCount === 0) {
    return ['no-usage',
      'no completions usage in the window, so there is nothing to attribute yet'];
  }

  const kinds = new Set(Object.keys(users).map((u) => classify(u, directory)));
  const principalNote = kinds.has('unresolved')
    ? `${userCount} user_id value(s), of which some resolve to nobody in the org directory`
    : `${userCount} user_id value(s), all of them org members or service accounts rather than customers`;

  if (keyCount === 1) {
    return ['single-key',
      '1 api_key_id covers every request in the window. There is one bucket, ' +
      `so per-customer cost has no place to come from. ${principalNote}.`];
  }

  if (tenantCount === null || tenantCount === undefined) {
    return ['unknown-tenant-count',
      `${keyCount} distinct api_key_id value(s) and ` +
      `${Object.keys(folded.projects ?? {}).length} project(s). ${principalNote}. ` +
      'Pass the tenant count to judge whether that is enough buckets.'];
  }

  if (keyCount < tenantCount) {
    return ['keys-below-tenants',
      `${keyCount} distinct api_key_id value(s) against ${tenantCount} ` +
      'tenant(s). Cost per customer is unrecoverable by construction: the ' +
      'finest slice the platform can offer is one key, and there are fewer ' +
      `keys than customers. ${principalNote}.`];
  }

  return ['segmented',
    `${keyCount} distinct api_key_id value(s) for ${tenantCount} tenant(s), so ` +
    'the platform can slice finely enough. Confirm your key-to-tenant map is current.'];
}

async function get(key, path, params = {}) {
  const url = new URL(API + path);
  for (const [k, v] of Object.entries(params)) {
    if (Array.isArray(v)) for (const item of v) url.searchParams.append(k, String(item));
    else if (v !== undefined && v !== null) url.searchParams.set(k, String(v));
  }
  const res = await fetch(url, { headers: { Authorization: `Bearer ${key}` } });
  if (res.status === 401) {
    throw new Error('401 from OpenAI: OPENAI_ADMIN_KEY must be an organization ' +
                    'admin key, not a project key');
  }
  if (res.status === 403) {
    throw new Error('403 from OpenAI: the key is not authorised for /v1/organization');
  }
  if (!res.ok) throw new Error(`${res.status} from ${path}`);
  return res.json();
}

async function usagePages(key, startTime, days, maxPages = 20) {
  const pages = [];
  let params = {
    start_time: startTime, bucket_width: '1d', limit: days, group_by: DIMENSIONS,
  };
  for (let i = 0; i < maxPages; i += 1) {
    const page = await get(key, '/organization/usage/completions', params);
    pages.push(page);
    if (!page.next_page) break;
    params = { ...params, page: page.next_page };
  }
  return pages;
}

async function orgDirectory(key, maxPages = 20) {
  const out = {};
  let params = { limit: 100 };
  for (let i = 0; i < maxPages; i += 1) {
    const page = await get(key, '/organization/users', params);
    const data = page.data ?? [];
    for (const user of data) {
      out[String(user.id)] = {
        name: user.name ?? user.email ?? '?',
        service_account: Boolean(user.is_service_account),
      };
    }
    if (!page.has_more || data.length === 0) break;
    params = { limit: 100, after: data[data.length - 1].id };
  }
  return out;
}

async function spendByKey(key, startTime) {
  const out = {};
  const page = await get(key, '/organization/costs',
    { start_time: startTime, limit: 30, group_by: 'api_key_id' });
  for (const bucket of page.data ?? []) {
    for (const result of bucket.results ?? []) {
      const keyId = String(result.api_key_id ?? 'unattributed');
      const amount = Number(result.amount?.value ?? 0);
      if (Number.isFinite(amount)) out[keyId] = (out[keyId] ?? 0) + amount;
    }
  }
  return out;
}

async function main() {
  const key = process.env.OPENAI_ADMIN_KEY;
  if (!key) {
    console.error('set OPENAI_ADMIN_KEY (an organization admin key with read scopes)');
    process.exitCode = 2;
    return;
  }

  const days = Number(process.env.DAYS ?? 7);
  const tenantsRaw = process.env.TENANTS;
  const tenants = tenantsRaw === undefined ? null : Number(tenantsRaw);

  const now = Math.floor(Date.now() / 1000);
  const folded = fold(await usagePages(key, now - days * 86400, days));
  const directory = await orgDirectory(key);
  const [state, detail] = verdict(folded, directory, tenants);

  console.log(`${state.padEnd(20)} ${detail}`);

  const byVolume = Object.keys(folded.users)
    .sort((a, b) => folded.users[b] - folded.users[a]);
  for (const userId of byVolume) {
    const kind = classify(userId, directory);
    const name = directory[userId]?.name ?? 'not in the directory';
    console.log(`  principal ${userId.padEnd(30)} ${kind.padEnd(16)} ${name}`);
  }

  const orphans = unresolved(folded, directory);
  if (orphans.length > 0) {
    console.warn(`  ${orphans.length} user_id value(s) resolve to nobody in the ` +
                 `org directory: ${orphans.join(', ')}`);
    console.warn('  repair: find what is calling as these principals before you ' +
                 'touch the attribution question');
  }

  if (FINDINGS.includes(state)) {
    const spend = await spendByKey(key, now - 30 * 86400);
    const top = Object.entries(spend).sort((a, b) => b[1] - a[1]).slice(0, 10);
    for (const [keyId, amount] of top) {
      console.warn(`  30d spend  ${keyId.padEnd(30)} $${amount.toFixed(2)}`);
    }
    console.warn('  repair: the Usage API cannot segment by end user. Mint one ' +
      'key, or one project, per tenant or tenant tier via ' +
      '/v1/organization/projects/{id}/service_accounts/{id}/api_keys and ' +
      'attribute with group_by=api_key_id.');
    console.warn('  repair: this is forward-only and cannot backfill. Until ' +
      'then, record each response usage block against your own tenant id and ' +
      'reconcile it against /v1/organization/costs.');
    process.exitCode = 1;
    return;
  }
  process.exitCode = 0;
}

// Only run when invoked directly, so importing this from the test file does not
// run main(), fail on the missing key, and set an exit code that fails the suite.
if (import.meta.url === `file://${process.argv[1]}`) {
  main().catch((err) => { console.error(err.message); process.exitCode = 2; });
}

from openai_tenant_attribution_audit import (classify, fold, unresolved,
                                             verdict)

DIRECTORY = {
    "user_eng1": {"name": "an engineer", "service_account": False},
    "user_eng2": {"name": "another engineer", "service_account": False},
    "sa_prod": {"name": "prod-backend", "service_account": True},
}


def folded(users=None, keys=None, projects=None, requests=100000):
    return {"users": users if users is not None else {"sa_prod": 100000},
            "keys": keys if keys is not None else {"key_abc": 100000},
            "projects": projects if projects is not None else {"proj_1": 100000},
            "requests": requests}


def bucket(rows):
    """One daily bucket from the usage endpoint, grouped three ways."""
    return {"data": [{"start_time": 0, "results": [
        {"user_id": u, "api_key_id": k, "project_id": p,
         "num_model_requests": n} for (u, k, p, n) in rows]}]}


def test_every_principal_is_one_of_your_own_and_that_is_the_finding():
    # Eleven rows, none of them a customer. The API answered; the answer is
    # about the org's own service accounts.
    state, detail = verdict(
        folded(users={"sa_prod": 90000, "user_eng1": 10000},
               keys={"key_a": 60000, "key_b": 40000}),
        DIRECTORY, tenant_count=412)
    assert state == "keys-below-tenants"
    assert "2 distinct api_key_id value(s) against 412 tenant(s)" in detail
    assert "org members or service accounts rather than customers" in detail


def test_one_key_is_its_own_worst_case():
    state, detail = verdict(folded(), DIRECTORY, tenant_count=412)
    assert state == "single-key"
    assert "one bucket" in detail


def test_enough_keys_means_the_platform_can_slice():
    state, _ = verdict(
        folded(keys={"key_%d" % i: 10 for i in range(500)}),
        DIRECTORY, tenant_count=412)
    assert state == "segmented"


def test_without_a_tenant_count_the_script_does_not_invent_a_verdict():
    state, detail = verdict(folded(keys={"key_a": 5, "key_b": 5}), DIRECTORY)
    assert state == "unknown-tenant-count"
    assert "Pass the tenant count" in detail
    assert verdict({"users": {}, "keys": {}, "projects": {}, "requests": 0},
                   DIRECTORY)[0] == "no-usage"


def test_a_principal_the_directory_does_not_know_is_a_different_problem():
    f = folded(users={"user_departed": 5000, "sa_prod": 5000},
               keys={"key_a": 5000, "key_b": 5000})
    assert classify("user_departed", DIRECTORY) == "unresolved"
    assert classify("sa_prod", DIRECTORY) == "service-account"
    assert classify("user_eng1", DIRECTORY) == "member"
    assert unresolved(f, DIRECTORY) == ["user_departed"]
    assert "resolve to nobody" in verdict(f, DIRECTORY, tenant_count=412)[1]


def test_fold_counts_the_three_dimensions_and_skips_the_nulls():
    pages = [bucket([("sa_prod", "key_a", "proj_1", 700),
                     (None, "key_b", "proj_1", 300)])]
    f = fold(pages)
    assert f["requests"] == 1000
    assert f["users"] == {"sa_prod": 700}
    assert f["keys"] == {"key_a": 700, "key_b": 300}
    assert f["projects"] == {"proj_1": 1000}

import { test } from 'node:test';
import assert from 'node:assert/strict';
import { classify, fold, unresolved, verdict }
  from './openai-tenant-attribution-audit.mjs';

const DIRECTORY = {
  user_eng1: { name: 'an engineer', service_account: false },
  user_eng2: { name: 'another engineer', service_account: false },
  sa_prod: { name: 'prod-backend', service_account: true },
};

function folded({ users, keys, projects, requests = 100000 } = {}) {
  return {
    users: users ?? { sa_prod: 100000 },
    keys: keys ?? { key_abc: 100000 },
    projects: projects ?? { proj_1: 100000 },
    requests,
  };
}

/** One daily bucket from the usage endpoint, grouped three ways. */
function bucket(rows) {
  return {
    data: [{
      start_time: 0,
      results: rows.map(([u, k, p, n]) => ({
        user_id: u, api_key_id: k, project_id: p, num_model_requests: n,
      })),
    }],
  };
}

test('every principal is one of your own and that is the finding', () => {
  const [state, detail] = verdict(
    folded({ users: { sa_prod: 90000, user_eng1: 10000 },
             keys: { key_a: 60000, key_b: 40000 } }),
    DIRECTORY, 412);
  assert.equal(state, 'keys-below-tenants');
  assert.match(detail, /2 distinct api_key_id value/);
  assert.match(detail, /org members or service accounts rather than customers/);
});

test('one key is its own worst case', () => {
  const [state, detail] = verdict(folded(), DIRECTORY, 412);
  assert.equal(state, 'single-key');
  assert.match(detail, /one bucket/);
});

test('enough keys means the platform can slice', () => {
  const keys = {};
  for (let i = 0; i < 500; i += 1) keys[`key_${i}`] = 10;
  assert.equal(verdict(folded({ keys }), DIRECTORY, 412)[0], 'segmented');
});

test('without a tenant count the script does not invent a verdict', () => {
  const [state, detail] = verdict(folded({ keys: { key_a: 5, key_b: 5 } }), DIRECTORY);
  assert.equal(state, 'unknown-tenant-count');
  assert.match(detail, /Pass the tenant count/);
  assert.equal(
    verdict({ users: {}, keys: {}, projects: {}, requests: 0 }, DIRECTORY)[0],
    'no-usage');
});

test('a principal the directory does not know is a different problem', () => {
  const f = folded({ users: { user_departed: 5000, sa_prod: 5000 },
                     keys: { key_a: 5000, key_b: 5000 } });
  assert.equal(classify('user_departed', DIRECTORY), 'unresolved');
  assert.equal(classify('sa_prod', DIRECTORY), 'service-account');
  assert.equal(classify('user_eng1', DIRECTORY), 'member');
  assert.deepEqual(unresolved(f, DIRECTORY), ['user_departed']);
  assert.match(verdict(f, DIRECTORY, 412)[1], /resolve to nobody/);
});

test('fold counts the three dimensions and skips the nulls', () => {
  const f = fold([bucket([['sa_prod', 'key_a', 'proj_1', 700],
                          [null, 'key_b', 'proj_1', 300]])]);
  assert.equal(f.requests, 1000);
  assert.deepEqual(f.users, { sa_prod: 700 });
  assert.deepEqual(f.keys, { key_a: 700, key_b: 300 });
  assert.deepEqual(f.projects, { proj_1: 1000 });
});