code:
/* Google Sheets backend for prompt logger.  Upon every flush, this backend
writes the batch to a Sheets spreadsheet. If the sheet becomes too large, it
will create a new sheet and continue writing there.

This is essentially a really shitty ORM for Sheets. Absolutely no concurrency
support because it relies on local state to match up with the remote state. */

import { google, sheets_v4 } from "googleapis";
import type { CredentialBody } from "google-auth-library";
import type { GaxiosResponse } from "googleapis-common";
import { config } from "../../../config";
import { logger } from "../../../logger";
import { PromptLogEntry } from "..";

// There is always a sheet called __index__ which contains a list of all the
// other sheets. We use this rather than iterating over all the sheets in case
// the user needs to manually work with the spreadsheet.
// If no __index__ sheet exists, we will assume that the spreadsheet is empty
// and create one.

type IndexSheetModel = {
  /**
   * Stored in cell B2. Set on startup; if it changes, we assume that another
   * instance of the proxy is writing to the spreadsheet and stop.
   */
  lockId: string;
  /**
   * Data starts at row 4. Row 1-3 are headers
   */
  rows: { logSheetName: string; createdAt: string; rowCount: number }[];
};

type LogSheetModel = {
  sheetName: string;
  rows: {
    model: string;
    endpoint: string;
    promptRaw: string;
    promptFlattened: string;
    response: string;
  }[];
};

const MAX_ROWS_PER_SHEET = 2000;
const log = logger.child({ module: "sheets" });

let sheetsClient: sheets_v4.Sheets | null = null;
/** Called when log backend aborts to tell the log queue to stop. */
let stopCallback: (() => void) | null = null;
/** Lock/synchronization ID for this session. */
let lockId = Math.random().toString(36).substring(2, 15);
/** In-memory cache of the index sheet. */
let indexSheet: IndexSheetModel | null = null;
/** In-memory cache of the active log sheet. */
let activeLogSheet: LogSheetModel | null = null;

/**
 * Loads the __index__ sheet into memory. By default, asserts that the lock ID
 * has not changed since the start of the session.
 */
const loadIndexSheet = async (assertLockId = true) => {
  const client = sheetsClient!;
  const spreadsheetId = config.googleSheetsSpreadsheetId!;
  log.info({ assertLockId }, "Loading __index__ sheet.");
  const res = await client.spreadsheets.values.get({
    spreadsheetId: spreadsheetId,
    range: "__index__!A1:D",
    majorDimension: "ROWS",
  });
  const data = assertData(res);
  if (!data.values || data.values[2][0] !== "logSheetName") {
    log.error({ values: data.values }, "Unexpected format for __index__ sheet");
    throw new Error("Unexpected format for __index__ sheet");
  }

  if (assertLockId) {
    const lockIdCell = data.values[1][1];
    if (lockIdCell !== lockId) {
      log.error(
        { receivedLock: lockIdCell, expectedLock: lockId },
        "Another instance of the proxy is writing to the spreadsheet; stopping."
      );
      stop();
      throw new Error(`Lock ID assertion failed`);
    }
  }

  const rows = data.values.slice(3).map((row) => {
    return {
      logSheetName: row[0],
      createdAt: row[1],
      rowCount: row[2],
    };
  });
  indexSheet = { lockId, rows };
};

/** Creates empty __index__ sheet for a new spreadsheet. */
const createIndexSheet = async () => {
  const client = sheetsClient!;
  const spreadsheetId = config.googleSheetsSpreadsheetId!;
  log.info("Creating empty __index__ sheet.");
  const res = await client.spreadsheets.batchUpdate({
    spreadsheetId: spreadsheetId,
    requestBody: {
      requests: [
        {
          addSheet: {
            properties: {
              title: "__index__",
              gridProperties: { rowCount: 1, columnCount: 3 },
            },
          },
        },
      ],
    },
  });
  assertData(res);
  indexSheet = { lockId, rows: [] };
  await writeIndexSheet();
};

/** Writes contents of in-memory indexSheet to the remote __index__ sheet. */
const writeIndexSheet = async () => {
  const client = sheetsClient!;
  const spreadsheetId = config.googleSheetsSpreadsheetId!;
  const headerRows = [
    ["Don't edit this sheet while the server is running.", "", ""],
    ["Lock ID", lockId, ""],
    ["logSheetName", "createdAt", "rowCount"],
  ];
  const contentRows = indexSheet!.rows.map((row) => {
    return [row.logSheetName, row.createdAt, row.rowCount];
  });
  log.info("Persisting __index__ sheet.");
  await client.spreadsheets.values.batchUpdate({
    spreadsheetId: spreadsheetId,
    requestBody: {
      valueInputOption: "RAW",
      data: [
        { range: "__index__!A1:D", values: [...headerRows, ...contentRows] },
      ],
    },
  });
};

/** Creates a new log sheet, adds it to the index, and sets it as active. */
const createLogSheet = async () => {
  const client = sheetsClient!;
  const spreadsheetId = config.googleSheetsSpreadsheetId!;
  // Sheet name format is Log_YYYYMMDD_HHMMSS
  const sheetName = `Log_${new Date()
    .toISOString()
    // YYYY-MM-DDTHH:MM:SS.sssZ -> YYYYMMDD_HHMMSS
    .replace(/[-:.]/g, "")
    .replace(/T/, "_")
    .substring(0, 15)}`;

  log.info({ sheetName }, "Creating new log sheet.");
  const res = await client.spreadsheets.batchUpdate({
    spreadsheetId: spreadsheetId,
    requestBody: {
      requests: [
        {
          addSheet: {
            properties: {
              title: sheetName,
              gridProperties: { rowCount: MAX_ROWS_PER_SHEET, columnCount: 5 },
            },
          },
        },
      ],
    },
  });
  assertData(res);
  // Increase row/column size and wrap text for readability.
  const sheetId = res.data.replies![0].addSheet!.properties!.sheetId;
  await client.spreadsheets.batchUpdate({
    spreadsheetId: spreadsheetId,
    requestBody: {
      requests: [
        {
          repeatCell: {
            range: { sheetId },
            cell: {
              userEnteredFormat: {
                wrapStrategy: "WRAP",
                verticalAlignment: "TOP",
              },
            },
            fields: "*",
          },
        },
        {
          updateDimensionProperties: {
            range: {
              sheetId,
              dimension: "COLUMNS",
              startIndex: 3,
              endIndex: 5,
            },
            properties: { pixelSize: 500 },
            fields: "pixelSize",
          },
        },
        {
          updateDimensionProperties: {
            range: {
              sheetId,
              dimension: "ROWS",
              startIndex: 1,
            },
            properties: { pixelSize: 200 },
            fields: "pixelSize",
          },
        },
      ],
    },
  });
  await client.spreadsheets.values.batchUpdate({
    spreadsheetId: spreadsheetId,
    requestBody: {
      valueInputOption: "RAW",
      data: [
        {
          range: `${sheetName}!A1:E`,
          values: [
            ["model", "endpoint", "prompt json", "prompt string", "response"],
          ],
        },
      ],
    },
  });
  indexSheet!.rows.push({
    logSheetName: sheetName,
    createdAt: new Date().toISOString(),
    rowCount: 0,
  });
  await writeIndexSheet();
  activeLogSheet = { sheetName, rows: [] };
};

export const appendBatch = async (batch: PromptLogEntry[]) => {
  if (!activeLogSheet) {
    // Create a new log sheet if we don't have one yet.
    await createLogSheet();
  } else {
    // Check lock to ensure we're the only instance writing to the spreadsheet.
    await loadIndexSheet(true);
  }

  const client = sheetsClient!;
  const spreadsheetId = config.googleSheetsSpreadsheetId!;
  const sheetName = activeLogSheet!.sheetName;
  const newRows = batch.map((entry) => {
    return [
      entry.model,
      entry.endpoint,
      entry.promptRaw.slice(-50000),
      entry.promptFlattened.slice(-50000),
      entry.response.slice(0, 50000),
    ];
  });
  log.info({ sheetName, rowCount: newRows.length }, "Appending log batch.");
  const data = await client.spreadsheets.values.append({
    spreadsheetId: spreadsheetId,
    range: `${sheetName}!A1:D`,
    valueInputOption: "RAW",
    requestBody: { values: newRows, majorDimension: "ROWS" },
  });
  assertData(data);
  if (data.data.updates && data.data.updates.updatedRows) {
    const newRowCount = data.data.updates.updatedRows;
    log.info({ sheetName, rowCount: newRowCount }, "Successfully appended.");
    activeLogSheet!.rows = activeLogSheet!.rows.concat(
      newRows.map((row) => ({
        model: row[0],
        endpoint: row[1],
        promptRaw: row[2],
        promptFlattened: row[3],
        response: row[4],
      }))
    );
  } else {
    // We didn't receive an error but we didn't get any updates either.
    // We may need to create a new sheet and throw to make the queue retry the
    // batch.
    log.warn(
      { sheetName, rowCount: newRows.length },
      "No updates received from append. Creating new sheet and retrying."
    );
    await createLogSheet();
    throw new Error("No updates received from append.");
  }
  await finalizeBatch();
};

const finalizeBatch = async () => {
  const sheetName = activeLogSheet!.sheetName;
  const rowCount = activeLogSheet!.rows.length;
  const indexRow = indexSheet!.rows.find(
    ({ logSheetName }) => logSheetName === sheetName
  )!;
  indexRow.rowCount = rowCount;
  if (rowCount >= MAX_ROWS_PER_SHEET) {
    await createLogSheet(); // Also updates index sheet
  } else {
    await writeIndexSheet();
  }
  log.info({ sheetName, rowCount }, "Batch finalized.");
};

type LoadLogSheetArgs = {
  sheetName: string;
  /** The starting row to load. If omitted, loads all rows (expensive). */
  fromRow?: number;
};

/** Not currently used. */
export const loadLogSheet = async ({
  sheetName,
  fromRow = 2, // omit header row
}: LoadLogSheetArgs) => {
  const client = sheetsClient!;
  const spreadsheetId = config.googleSheetsSpreadsheetId!;

  const range = `${sheetName}!A${fromRow}:E`;
  const res = await client.spreadsheets.values.get({
    spreadsheetId: spreadsheetId,
    range,
  });
  const data = assertData(res);
  const values = data.values || [];
  const rows = values.slice(1).map((row) => {
    return {
      model: row[0],
      endpoint: row[1],
      promptRaw: row[2],
      promptFlattened: row[3],
      response: row[4],
    };
  });
  activeLogSheet = { sheetName, rows };
};

export const init = async (onStop: () => void) => {
  if (sheetsClient) {
    return;
  }
  if (!config.googleSheetsKey || !config.googleSheetsSpreadsheetId) {
    throw new Error(
      "Missing required Google Sheets config. Refer to documentation for setup instructions."
    );
  }

  log.info("Initializing Google Sheets backend.");
  const encodedCreds = config.googleSheetsKey;
  // encodedCreds is a base64-encoded JSON key from the GCP console.
  const creds: CredentialBody = JSON.parse(
    Buffer.from(encodedCreds, "base64").toString("utf8").trim()
  );
  const auth = new google.auth.GoogleAuth({
    scopes: ["https://www.googleapis.com/auth/spreadsheets"],
    credentials: creds,
  });
  sheetsClient = google.sheets({ version: "v4", auth });
  stopCallback = onStop;

  const sheetId = config.googleSheetsSpreadsheetId;
  const res = await sheetsClient.spreadsheets.get({
    spreadsheetId: sheetId,
  });
  if (!res.data) {
    const { status, statusText, headers } = res;
    log.error(
      {
        res: { status, statusText, headers },
        creds: {
          client_email: creds.client_email?.slice(0, 5) + "********",
          private_key: creds.private_key?.slice(0, 5) + "********",
        },
        sheetId: config.googleSheetsSpreadsheetId,
      },
      "Could not connect to Google Sheets."
    );
    stop();
    throw new Error("Could not connect to Google Sheets.");
  } else {
    const sheetTitle = res.data.properties?.title;
    log.info({ sheetId, sheetTitle }, "Connected to Google Sheets.");
  }

  // Load or create the index sheet and write the lockId to it.
  try {
    log.info("Loading index sheet.");
    await loadIndexSheet(false);
    await writeIndexSheet();
  } catch (e) {
    log.warn({ error: e.message }, "Could not load index sheet. Creating a new one.");
    await createIndexSheet();
  }
};

/** Called during some unrecoverable error to tell the log queue to stop. */
function stop() {
  log.warn("Stopping Google Sheets backend.");
  if (stopCallback) {
    stopCallback();
  }
  sheetsClient = null;
}

function assertData<T = sheets_v4.Schema$ValueRange>(res: GaxiosResponse<T>) {
  if (!res.data) {
    const { status, statusText, headers } = res;
    log.error(
      { res: { status, statusText, headers } },
      "Unexpected response from Google Sheets API."
    );
  }
  return res.data!;
}


code:
/* Logs prompts and model responses to a persistent storage backend, if enabled.
Since the proxy is generally deployed to free-tier services, our options for
persistent storage are pretty limited. We'll use Google Sheets as a makeshift
database for now.

Due to the limitations of Google Sheets, we'll queue up log entries and flush
them to the API periodically. */

export interface PromptLogEntry {
  model: string;
  endpoint: string;
  /** JSON prompt passed to the model */
  promptRaw: string;
  /** Prompt with user and assistant messages flattened into a single string */
  promptFlattened: string;
  response: string;
  // TODO: temperature, top_p, top_k, etc.
}

export * as logQueue from "./log-queue";

repo_name: Cucumber/oai

license: unknown


code:
/* Queues incoming prompts/responses and periodically flushes them to configured
 * logging backend. */

import { logger } from "../../logger";
import { PromptLogEntry } from ".";
import { sheets } from "./backends";

const FLUSH_INTERVAL = 1000 * 10;
const MAX_BATCH_SIZE = 25;

const queue: PromptLogEntry[] = [];
const log = logger.child({ module: "log-queue" });

let started = false;
let timeoutId: NodeJS.Timeout | null = null;
let retrying = false;
let consecutiveFailedBatches = 0;

export const enqueue = (payload: PromptLogEntry) => {
  if (!started) {
    log.warn("Log queue not started, discarding incoming log entry.");
    return;
  }
  queue.push(payload);
};

export const flush = async () => {
  if (!started) {
    return;
  }

  if (queue.length > 0) {
    const batchSize = Math.min(MAX_BATCH_SIZE, queue.length);
    const nextBatch = queue.splice(0, batchSize);
    log.info({ size: nextBatch.length }, "Submitting new batch.");
    try {
      await sheets.appendBatch(nextBatch);
      retrying = false;
      consecutiveFailedBatches = 0;
    } catch (e: any) {
      if (retrying) {
        log.error(
          { message: e.message, stack: e.stack },
          "Failed twice to flush batch, discarding."
        );
        retrying = false;
        consecutiveFailedBatches++;
      } else {
        // Put the batch back at the front of the queue and try again
        log.warn(
          { message: e.message, stack: e.stack },
          "Failed to flush batch. Retrying."
        );
        queue.unshift(...nextBatch);
        retrying = true;
        setImmediate(() => flush());
        return;
      }
    }
  }

  const useHalfInterval = queue.length > MAX_BATCH_SIZE / 2;
  scheduleFlush(useHalfInterval);
};

export const start = async () => {
  try {
    await sheets.init(() => stop());
    log.info("Logging backend initialized.");
    started = true;
  } catch (e) {
    log.error({ error: e.message }, "Could not initialize logging backend.");
    return;
  }
  scheduleFlush();
};

export const stop = () => {
  if (timeoutId) {
    clearTimeout(timeoutId);
  }
  log.info("Stopping log queue.");
  started = false;
};

const scheduleFlush = (halfInterval = false) => {
  if (consecutiveFailedBatches > 3) {
    // TODO: may cause memory issues on busy servers, though if we crash that
    // may actually fix the problem with logs randomly not being flushed.
    const oneMinute = 60 * 1000;
    const maxBackoff = 10 * oneMinute;
    const backoff = Math.min(consecutiveFailedBatches * oneMinute, maxBackoff);
    timeoutId = setTimeout(() => {
      flush();
    }, backoff);
    log.warn(
      { consecutiveFailedBatches, backoffMs: backoff },
      "Failed to flush 3 batches in a row, pausing for a few minutes."
    );
    return;
  }

  if (halfInterval) {
    log.warn(
      { queueSize: queue.length },
      "Queue is falling behind, switching to faster flush interval."
    );
  }

  timeoutId = setTimeout(
    () => {
      flush();
    },
    halfInterval ? FLUSH_INTERVAL / 2 : FLUSH_INTERVAL
  );
};


code:
import { config } from "../config";
import { ModelFamily } from "./models";

// technically slightly underestimates, because completion tokens cost more
// than prompt tokens but we don't track those separately right now
export function getTokenCostUsd(model: ModelFamily, tokens: number) {
  let cost = 0;
  switch (model) {
    case "azure-gpt4-turbo":
    case "gpt4-turbo":
      cost = 0.00001;
      break;
    case "azure-gpt4-32k":
    case "gpt4-32k":
      cost = 0.00006;
      break;
    case "azure-gpt4":
    case "gpt4":
      cost = 0.00003;
      break;
    case "azure-turbo":
    case "turbo":
      cost = 0.000001;
      break;
    case "dall-e":
      cost = 0.00001;
      break;
    case "aws-claude":
    case "claude":
      cost = 0.000008;
      break;
    case "claude-opus":
      cost = 0.000015;
      break;
    case "mistral-tiny":
      cost = 0.00000031;
      break;
    case "mistral-small":
      cost = 0.00000132;
      break;
    case "mistral-medium":
      cost = 0.0000055;
      break;
  }
  return cost * Math.max(0, tokens);
}

export function prettyTokens(tokens: number): string {
  const absTokens = Math.abs(tokens);
  if (absTokens < 1000) {
    return tokens.toString();
  } else if (absTokens < 1000000) {
    return (tokens / 1000).toFixed(1) + "k";
  } else if (absTokens < 1000000000) {
    return (tokens / 1000000).toFixed(2) + "m";
  } else {
    return (tokens / 1000000000).toFixed(3) + "b";
  }
}

export function getCostSuffix(cost: number) {
  if (!config.showTokenCosts) return "";
  return ` ($${cost.toFixed(2)})`;
}


code:
import { Response } from "express";
import { IncomingMessage } from "http";

export function initializeSseStream(res: Response) {
  res.statusCode = 200;
  res.setHeader("Content-Type", "text/event-stream; charset=utf-8");
  res.setHeader("Cache-Control", "no-cache");
  res.setHeader("Connection", "keep-alive");
  res.setHeader("X-Accel-Buffering", "no"); // nginx-specific fix
  res.flushHeaders();
}

/**
 * Copies headers received from upstream API to the SSE response, excluding
 * ones we need to set ourselves for SSE to work.
 */
export function copySseResponseHeaders(
  proxyRes: IncomingMessage,
  res: Response
) {
  const toOmit = [
    "content-length",
    "content-encoding",
    "transfer-encoding",
    "content-type",
    "connection",
    "cache-control",
  ];
  for (const [key, value] of Object.entries(proxyRes.headers)) {
    if (!toOmit.includes(key) && value) {
      res.setHeader(key, value);
    }
  }
}


code:
import { getTokenizer } from "@anthropic-ai/tokenizer";
import { Tiktoken } from "tiktoken/lite";
import { AnthropicChatMessage } from "../api-schemas";
import { libSharp } from "../file-storage";
import { logger } from "../../logger";

const log = logger.child({ module: "tokenizer", service: "anthropic" });

let encoder: Tiktoken;
let userRoleCount = 0;
let assistantRoleCount = 0;

export function init() {
  // they export a `countTokens` function too but it instantiates a new
  // tokenizer every single time and it is not fast...
  encoder = getTokenizer();
  userRoleCount = encoder.encode("\n\nHuman: ", "all").length;
  assistantRoleCount = encoder.encode("\n\nAssistant: ", "all").length;
  return true;
}

export async function getTokenCount(prompt: string | AnthropicChatMessage[]) {
  if (typeof prompt !== "string") {
    return getTokenCountForMessages(prompt);
  }

  if (prompt.length > 800000) {
    throw new Error("Content is too large to tokenize.");
  }

  return {
    tokenizer: "@anthropic-ai/tokenizer",
    token_count: encoder.encode(prompt.normalize("NFKC"), "all").length,
  };
}

async function getTokenCountForMessages(messages: AnthropicChatMessage[]) {
  let numTokens = 0;

  for (const message of messages) {
    const { content, role } = message;
    numTokens += role === "user" ? userRoleCount : assistantRoleCount;

    const parts = Array.isArray(content)
      ? content
      : [{ type: "text" as const, text: content }];

    for (const part of parts) {
      switch (part.type) {
        case "text":
          const { text } = part;
          if (text.length > 800000 || numTokens > 200000) {
            throw new Error("Text content is too large to tokenize.");
          }
          numTokens += encoder.encode(text.normalize("NFKC"), "all").length;
          break;
        case "image":
          numTokens += await getImageTokenCount(part.source.data);
          break;
        default:
          throw new Error(`Unsupported Anthropic content type.`);
      }
    }
  }

  if (messages[messages.length - 1].role !== "assistant") {
    numTokens += assistantRoleCount;
  }

  return { tokenizer: "@anthropic-ai/tokenizer", token_count: numTokens };
}

async function getImageTokenCount(b64: string) {
  // https://docs.anthropic.com/claude/docs/vision
  // If your image's long edge is more than 1568 pixels, or your image is more
  // than ~1600 tokens, it will first be scaled down, preserving aspect ratio,
  // until it is within size limits. Assuming your image does not need to be
  // resized, you can estimate the number of tokens used via this simple
  // algorithm:
  // tokens = (width px * height px)/750

  const buffer = Buffer.from(b64, "base64");
  const image = libSharp(buffer);
  const metadata = await image.metadata();

  if (!metadata || !metadata.width || !metadata.height) {
    throw new Error("Prompt includes an image that could not be parsed");
  }

  const MAX_TOKENS = 1600;
  const MAX_LENGTH_PX = 1568;
  const PIXELS_PER_TOKEN = 750;
  const { width, height } = metadata;
  let tokens = (width * height) / PIXELS_PER_TOKEN;

  // Resize the image if it's too large
  if (tokens > MAX_TOKENS || width > MAX_LENGTH_PX || height > MAX_LENGTH_PX) {
    const longestEdge = Math.max(width, height);

    let factor;
    if (tokens > MAX_TOKENS) {
      const targetPixels = PIXELS_PER_TOKEN * MAX_TOKENS;
      factor = Math.sqrt(targetPixels / (width * height));
    } else {
      factor = MAX_LENGTH_PX / longestEdge;
    }

    const scaledWidth = width * factor;
    const scaledHeight = height * factor;

    tokens = (scaledWidth * scaledHeight) / 750;
  }

  log.debug({ width, height, tokens }, "Calculated Claude Vision token cost");
  return Math.ceil(tokens);
}


code: export { init, countTokens } from "./tokenizer";


code:
import * as tokenizer from "./mistral-tokenizer-js";
import { MistralAIChatMessage } from "../api-schemas";

export function init() {
  tokenizer.initializemistralTokenizer();
  return true;
}

export function getTokenCount(prompt: MistralAIChatMessage[] | string) {
  if (typeof prompt === "string") {
    return getTextTokenCount(prompt);
  }

  let chunks = [];
  for (const message of prompt) {
    switch (message.role) {
      case "system":
        chunks.push(message.content);
        break;
      case "assistant":
        chunks.push(message.content + "</s>");
        break;
      case "user":
        chunks.push("[INST] " + message.content + " [/INST]");
        break;
    }
  }
  return getTextTokenCount(chunks.join(" "));
}

function getTextTokenCount(prompt: string) {
  if (prompt.length > 800000) {
    throw new Error("Content is too large to tokenize.");
  }

  return {
    tokenizer: "mistral-tokenizer-js",
    token_count: tokenizer.encode(prompt.normalize("NFKC"))!.length,
  };
}


code:
import { Tiktoken } from "tiktoken/lite";
import cl100k_base from "tiktoken/encoders/cl100k_base.json";
import { logger } from "../../logger";
import { libSharp } from "../file-storage";
import { GoogleAIChatMessage, OpenAIChatMessage } from "../api-schemas";

const log = logger.child({ module: "tokenizer", service: "openai" });
const GPT4_VISION_SYSTEM_PROMPT_SIZE = 170;

let encoder: Tiktoken;

export function init() {
  encoder = new Tiktoken(
    cl100k_base.bpe_ranks,
    cl100k_base.special_tokens,
    cl100k_base.pat_str
  );
  return true;
}

// Tested against:
// https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb

export async function getTokenCount(
  prompt: string | OpenAIChatMessage[],
  model: string
) {
  if (typeof prompt === "string") {
    return getTextTokenCount(prompt);
  }

  const oldFormatting = model.startsWith("turbo-0301");
  const vision = model.includes("vision");

  const tokensPerMessage = oldFormatting ? 4 : 3;
  const tokensPerName = oldFormatting ? -1 : 1; // older formatting replaces role with name if name is present

  let numTokens = vision ? GPT4_VISION_SYSTEM_PROMPT_SIZE : 0;

  for (const message of prompt) {
    numTokens += tokensPerMessage;
    for (const key of Object.keys(message)) {
      {
        let textContent: string = "";
        const value = message[key as keyof OpenAIChatMessage];

        if (!value) continue;

        if (Array.isArray(value)) {
          for (const item of value) {
            if (item.type === "text") {
              textContent += item.text;
            } else if (["image", "image_url"].includes(item.type)) {
              const { url, detail } = item.image_url;
              const cost = await getGpt4VisionTokenCost(url, detail);
              numTokens += cost ?? 0;
            }
          }
        } else {
          textContent = value;
        }

        if (textContent.length > 800000 || numTokens > 200000) {
          throw new Error("Content is too large to tokenize.");
        }

        numTokens += encoder.encode(textContent).length;
        if (key === "name") {
          numTokens += tokensPerName;
        }
      }
    }
  }
  numTokens += 3; // every reply is primed with <|start|>assistant<|message|>
  return { tokenizer: "tiktoken", token_count: numTokens };
}

async function getGpt4VisionTokenCost(
  url: string,
  detail: "auto" | "low" | "high" = "auto"
) {
  // For now we do not allow remote images as the proxy would have to download
  // them, which is a potential DoS vector.
  if (!url.startsWith("data:image/")) {
    throw new Error(
      "Remote images are not supported. Add the image to your prompt as a base64 data URL."
    );
  }

  const base64Data = url.split(",")[1];
  const buffer = Buffer.from(base64Data, "base64");
  const image = libSharp(buffer);
  const metadata = await image.metadata();

  if (!metadata || !metadata.width || !metadata.height) {
    throw new Error("Prompt includes an image that could not be parsed");
  }

  const { width, height } = metadata;

  let selectedDetail: "low" | "high";
  if (detail === "auto") {
    const threshold = 512 * 512;
    const imageSize = width * height;
    selectedDetail = imageSize > threshold ? "high" : "low";
  } else {
    selectedDetail = detail;
  }

  // https://platform.openai.com/docs/guides/vision/calculating-costs
  if (selectedDetail === "low") {
    log.info(
      { width, height, tokens: 85 },
      "Using fixed GPT-4-Vision token cost for low detail image"
    );
    return 85;
  }

  let newWidth = width;
  let newHeight = height;
  if (width > 2048 || height > 2048) {
    const aspectRatio = width / height;
    if (width > height) {
      newWidth = 2048;
      newHeight = Math.round(2048 / aspectRatio);
    } else {
      newHeight = 2048;
      newWidth = Math.round(2048 * aspectRatio);
    }
  }

  if (newWidth < newHeight) {
    newHeight = Math.round((newHeight / newWidth) * 768);
    newWidth = 768;
  } else {
    newWidth = Math.round((newWidth / newHeight) * 768);
    newHeight = 768;
  }

  const tiles = Math.ceil(newWidth / 512) * Math.ceil(newHeight / 512);
  const tokens = 170 * tiles + 85;

  log.info(
    { width, height, newWidth, newHeight, tiles, tokens },
    "Calculated GPT-4-Vision token cost for high detail image"
  );
  return tokens;
}

function getTextTokenCount(prompt: string) {
  if (prompt.length > 500000) {
    return {
      tokenizer: "length fallback",
      token_count: 100000,
    };
  }

  return {
    tokenizer: "tiktoken",
    token_count: encoder.encode(prompt).length,
  };
}

// Model	Resolution	Price
// DALL·E 3	1024×1024	$0.040 / image
// 1024×1792, 1792×1024	$0.080 / image
// DALL·E 3 HD	1024×1024	$0.080 / image
// 1024×1792, 1792×1024	$0.120 / image
// DALL·E 2	1024×1024	$0.020 / image
// 512×512	$0.018 / image
// 256×256	$0.016 / image

export const DALLE_TOKENS_PER_DOLLAR = 100000;

/**
 * OpenAI image generation with DALL-E doesn't use tokens but everything else
 * in the application does. There is a fixed cost for each image generation
 * request depending on the model and selected quality/resolution parameters,
 * which we convert to tokens at a rate of 100000 tokens per dollar.
 */
export function getOpenAIImageCost(params: {
  model: "dall-e-2" | "dall-e-3";
  quality: "standard" | "hd";
  resolution: "512x512" | "256x256" | "1024x1024" | "1024x1792" | "1792x1024";
  n: number | null;
}) {
  const { model, quality, resolution, n } = params;
  const usd = (() => {
    switch (model) {
      case "dall-e-2":
        switch (resolution) {
          case "512x512":
            return 0.018;
          case "256x256":
            return 0.016;
          case "1024x1024":
            return 0.02;
          default:
            throw new Error("Invalid resolution");
        }
      case "dall-e-3":
        switch (resolution) {
          case "1024x1024":
            return quality === "standard" ? 0.04 : 0.08;
          case "1024x1792":
          case "1792x1024":
            return quality === "standard" ? 0.08 : 0.12;
          default:
            throw new Error("Invalid resolution");
        }
      default:
        throw new Error("Invalid image generation model");
    }
  })();

  const tokens = (n ?? 1) * (usd * DALLE_TOKENS_PER_DOLLAR);

  return {
    tokenizer: `openai-image cost`,
    token_count: Math.ceil(tokens),
  };
}

export function estimateGoogleAITokenCount(
  prompt: string | GoogleAIChatMessage[]
) {
  if (typeof prompt === "string") {
    return getTextTokenCount(prompt);
  }

  const tokensPerMessage = 3;

  let numTokens = 0;
  for (const message of prompt) {
    numTokens += tokensPerMessage;
    numTokens += encoder.encode(message.parts[0].text).length;
  }

  numTokens += 3;

  return {
    tokenizer: "tiktoken (google-ai estimate)",
    token_count: numTokens,
  };
}


code:
import { Request } from "express";
import { assertNever } from "../utils";
import {
  getTokenCount as getClaudeTokenCount,
  init as initClaude,
} from "./claude";
import {
  estimateGoogleAITokenCount,
  getOpenAIImageCost,
  getTokenCount as getOpenAITokenCount,
  init as initOpenAi,
} from "./openai";
import {
  getTokenCount as getMistralAITokenCount,
  init as initMistralAI,
} from "./mistral";
import { APIFormat } from "../key-management";
import {
  AnthropicChatMessage,
  GoogleAIChatMessage,
  MistralAIChatMessage,
  OpenAIChatMessage,
} from "../api-schemas";

export async function init() {
  initClaude();
  initOpenAi();
  initMistralAI();
}

type OpenAIChatTokenCountRequest = {
  prompt: OpenAIChatMessage[];
  completion?: never;
  service: "openai";
};

type AnthropicChatTokenCountRequest = {
  prompt: AnthropicChatMessage[];
  completion?: never;
  service: "anthropic-chat";
};

type GoogleAIChatTokenCountRequest = {
  prompt: GoogleAIChatMessage[];
  completion?: never;
  service: "google-ai";
};

type MistralAIChatTokenCountRequest = {
  prompt: MistralAIChatMessage[];
  completion?: never;
  service: "mistral-ai";
};

type FlatPromptTokenCountRequest = {
  prompt: string;
  completion?: never;
  service: "openai-text" | "anthropic-text" | "google-ai";
};

type StringCompletionTokenCountRequest = {
  prompt?: never;
  completion: string;
  service: APIFormat;
};

type OpenAIImageCompletionTokenCountRequest = {
  prompt?: never;
  completion?: never;
  service: "openai-image";
};

/**
 * Tagged union via `service` field of the different types of requests that can
 * be made to the tokenization service, for both prompts and completions
 */
type TokenCountRequest = { req: Request } & (
  | OpenAIChatTokenCountRequest
  | AnthropicChatTokenCountRequest
  | GoogleAIChatTokenCountRequest
  | MistralAIChatTokenCountRequest
  | FlatPromptTokenCountRequest
  | StringCompletionTokenCountRequest
  | OpenAIImageCompletionTokenCountRequest
);

type TokenCountResult = {
  token_count: number;
  tokenizer: string;
  tokenization_duration_ms: number;
};

export async function countTokens({
  req,
  service,
  prompt,
  completion,
}: TokenCountRequest): Promise<TokenCountResult> {
  const time = process.hrtime();
  switch (service) {
    case "anthropic-chat":
    case "anthropic-text":
      return {
        ...(await getClaudeTokenCount(prompt ?? completion)),
        tokenization_duration_ms: getElapsedMs(time),
      };
    case "openai":
    case "openai-text":
      return {
        ...(await getOpenAITokenCount(prompt ?? completion, req.body.model)),
        tokenization_duration_ms: getElapsedMs(time),
      };
    case "openai-image":
      return {
        ...getOpenAIImageCost({
          model: req.body.model,
          quality: req.body.quality,
          resolution: req.body.size,
          n: parseInt(req.body.n, 10) || null,
        }),
        tokenization_duration_ms: getElapsedMs(time),
      };
    case "google-ai":
      // TODO: Can't find a tokenization library for Gemini. There is an API
      // endpoint for it but it adds significant latency to the request.
      return {
        ...estimateGoogleAITokenCount(prompt ?? (completion || [])),
        tokenization_duration_ms: getElapsedMs(time),
      };
    case "mistral-ai":
      return {
        ...getMistralAITokenCount(prompt ?? completion),
        tokenization_duration_ms: getElapsedMs(time),
      };
    default:
      assertNever(service);
  }
}

function getElapsedMs(time: [number, number]) {
  const diff = process.hrtime(time);
  return diff[0] * 1000 + diff[1] / 1e6;
}


code:
import { ZodType, z } from "zod";
import type { ModelFamily } from "../models";
import { makeOptionalPropsNullable } from "../utils";

export const tokenCountsSchema: ZodType<UserTokenCounts> = z.object({
  turbo: z.number().optional().default(0),
  gpt4: z.number().optional().default(0),
  "gpt4-32k": z.number().optional().default(0),
  "gpt4-turbo": z.number().optional().default(0),
  "dall-e": z.number().optional().default(0),
  claude: z.number().optional().default(0),
  "gemini-pro": z.number().optional().default(0),
  "aws-claude": z.number().optional().default(0),
});

export const UserSchema = z
  .object({
    /** User's personal access token. */
    token: z.string(),
    /** IP addresses the user has connected from. */
    ip: z.array(z.string()),
    /** User's nickname. */
    nickname: z.string().max(80).optional(),
    /**
     * The user's privilege level.
     * - `normal`: Default role. Subject to usual rate limits and quotas.
     * - `special`: Special role. Higher quotas and exempt from
     *   auto-ban/lockout.
     **/
    type: z.enum(["normal", "special", "temporary"]),
    /** Number of prompts the user has made. */
    promptCount: z.number(),
    /**
     * @deprecated Use `tokenCounts` instead.
     * Never used; retained for backwards compatibility.
     */
    tokenCount: z.any().optional(),
    /** Number of tokens the user has consumed, by model family. */
    tokenCounts: tokenCountsSchema,
    /** Maximum number of tokens the user can consume, by model family. */
    tokenLimits: tokenCountsSchema,
    /** Time at which the user was created. */
    createdAt: z.number(),
    /** Time at which the user last connected. */
    lastUsedAt: z.number().optional(),
    /** Time at which the user was disabled, if applicable. */
    disabledAt: z.number().optional(),
    /** Reason for which the user was disabled, if applicable. */
    disabledReason: z.string().optional(),
    /** Time at which the user will expire and be disabled (for temp users). */
    expiresAt: z.number().optional(),
    /** The user's maximum number of IP addresses; supercedes global max. */
    maxIps: z.coerce.number().int().min(0).optional(),
    /** Private note about the user. */
    adminNote: z.string().optional(),
  })
  .strict();

/**
 * Variant of `UserSchema` which allows for partial updates, and makes any
 * optional properties on the base schema nullable. Null values are used to
 * indicate that the property should be deleted from the user object.
 */
export const UserPartialSchema = makeOptionalPropsNullable(UserSchema)
  .partial()
  .extend({ token: z.string() });

export type UserTokenCounts = {
  [K in ModelFamily]?: number;
};
export type User = z.infer<typeof UserSchema>;
export type UserUpdate = z.infer<typeof UserPartialSchema>;


code:
/**
 * Basic user management. Handles creation and tracking of proxy users, personal
 * access tokens, and quota management. Supports in-memory and Firebase Realtime
 * Database persistence stores.
 *
 * Users are identified solely by their personal access token. The token is
 * used to authenticate the user for all proxied requests.
 */

import admin from "firebase-admin";
import schedule from "node-schedule";
import { v4 as uuid } from "uuid";
import { config, getFirebaseApp } from "../../config";
import {
  getAwsBedrockModelFamily,
  getAzureOpenAIModelFamily,
  getClaudeModelFamily,
  getGoogleAIModelFamily,
  getMistralAIModelFamily,
  getOpenAIModelFamily,
  MODEL_FAMILIES,
  ModelFamily,
} from "../models";
import { logger } from "../../logger";
import { User, UserTokenCounts, UserUpdate } from "./schema";
import { APIFormat } from "../key-management";
import { assertNever } from "../utils";

const log = logger.child({ module: "users" });

const INITIAL_TOKENS: Required<UserTokenCounts> = {
  turbo: 0,
  gpt4: 0,
  "gpt4-32k": 0,
  "gpt4-turbo": 0,
  "dall-e": 0,
  claude: 0,
  "claude-opus": 0,
  "gemini-pro": 0,
  "mistral-tiny": 0,
  "mistral-small": 0,
  "mistral-medium": 0,
  "mistral-large": 0,
  "aws-claude": 0,
  "azure-turbo": 0,
  "azure-gpt4": 0,
  "azure-gpt4-turbo": 0,
  "azure-gpt4-32k": 0,
};

const users: Map<string, User> = new Map();
const usersToFlush = new Set<string>();
let quotaRefreshJob: schedule.Job | null = null;
let userCleanupJob: schedule.Job | null = null;

export async function init() {
  log.info({ store: config.gatekeeperStore }, "Initializing user store...");
  if (config.gatekeeperStore === "firebase_rtdb") {
    await initFirebase();
  }
  if (config.quotaRefreshPeriod) {
    const crontab = getRefreshCrontab();
    quotaRefreshJob = schedule.scheduleJob(crontab, refreshAllQuotas);
    if (!quotaRefreshJob) {
      throw new Error(
        "Unable to schedule quota refresh. Is QUOTA_REFRESH_PERIOD set correctly?"
      );
    }
    log.debug(
      { nextRefresh: quotaRefreshJob.nextInvocation() },
      "Scheduled token quota refresh."
    );
  }

  userCleanupJob = schedule.scheduleJob("* * * * *", cleanupExpiredTokens);

  log.info("User store initialized.");
}

/**
 * Creates a new user and returns their token. Optionally accepts parameters
 * for setting an expiry date and/or token limits for temporary users.
 **/
export function createUser(createOptions?: {
  type?: User["type"];
  expiresAt?: number;
  tokenLimits?: User["tokenLimits"];
}) {
  const token = uuid();
  const newUser: User = {
    token,
    ip: [],
    type: "normal",
    promptCount: 0,
    tokenCounts: { ...INITIAL_TOKENS },
    tokenLimits: createOptions?.tokenLimits ?? { ...config.tokenQuota },
    createdAt: Date.now(),
  };

  if (createOptions?.type === "temporary") {
    Object.assign(newUser, {
      type: "temporary",
      expiresAt: createOptions.expiresAt,
    });
  } else {
    Object.assign(newUser, { type: createOptions?.type ?? "normal" });
  }

  users.set(token, newUser);
  usersToFlush.add(token);
  return token;
}

/** Returns the user with the given token if they exist. */
export function getUser(token: string) {
  return users.get(token);
}

/** Returns a list of all users. */
export function getUsers() {
  return Array.from(users.values()).map((user) => ({ ...user }));
}

/**
 * Upserts the given user. Intended for use with the /admin API for updating
 * arbitrary fields on a user; use the other functions in this module for
 * specific use cases. `undefined` values are left unchanged. `null` will delete
 * the property from the user.
 *
 * Returns the upserted user.
 */
export function upsertUser(user: UserUpdate) {
  const existing: User = users.get(user.token) ?? {
    token: user.token,
    ip: [],
    type: "normal",
    promptCount: 0,
    tokenCounts: { ...INITIAL_TOKENS },
    tokenLimits: { ...config.tokenQuota },
    createdAt: Date.now(),
  };

  const updates: Partial<User> = {};

  for (const field of Object.entries(user)) {
    const [key, value] = field as [keyof User, any]; // already validated by zod
    if (value === undefined || key === "token") continue;
    if (value === null) {
      delete existing[key];
    } else {
      updates[key] = value;
    }
  }

  // TODO: Write firebase migration to backfill new fields
  if (updates.tokenCounts) {
    for (const family of MODEL_FAMILIES) {
      updates.tokenCounts[family] ??= 0;
    }
  }
  if (updates.tokenLimits) {
    for (const family of MODEL_FAMILIES) {
      updates.tokenLimits[family] ??= 0;
    }
  }

  users.set(user.token, Object.assign(existing, updates));
  usersToFlush.add(user.token);

  // Immediately schedule a flush to the database if we're using Firebase.
  if (config.gatekeeperStore === "firebase_rtdb") {
    setImmediate(flushUsers);
  }

  return users.get(user.token);
}

/** Increments the prompt count for the given user. */
export function incrementPromptCount(token: string) {
  const user = users.get(token);
  if (!user) return;
  user.promptCount++;
  usersToFlush.add(token);
}

/** Increments token consumption for the given user and model. */
export function incrementTokenCount(
  token: string,
  model: string,
  api: APIFormat,
  consumption: number
) {
  const user = users.get(token);
  if (!user) return;
  const modelFamily = getModelFamilyForQuotaUsage(model, api);
  const existing = user.tokenCounts[modelFamily] ?? 0;
  user.tokenCounts[modelFamily] = existing + consumption;
  usersToFlush.add(token);
}

/**
 * Given a user's token and IP address, authenticates the user and adds the IP
 * to the user's list of IPs. Returns the user if they exist and are not
 * disabled, otherwise returns undefined.
 */
export function authenticate(
  token: string,
  ip: string
): { user?: User; result: "success" | "disabled" | "not_found" | "limited" } {
  const user = users.get(token);
  if (!user) return { result: "not_found" };
  if (user.disabledAt) return { result: "disabled" };

  const newIp = !user.ip.includes(ip);

  const userLimit = user.maxIps ?? config.maxIpsPerUser;
  const enforcedLimit =
    user.type === "special" || !userLimit ? Infinity : userLimit;

  if (newIp && user.ip.length >= enforcedLimit) {
    if (config.maxIpsAutoBan) {
      user.ip.push(ip);
      disableUser(token, "IP address limit exceeded.");
      return { result: "disabled" };
    }
    return { result: "limited" };
  } else if (newIp) {
    user.ip.push(ip);
  }

  user.lastUsedAt = Date.now();
  usersToFlush.add(token);
  return { user, result: "success" };
}

export function hasAvailableQuota({
  userToken,
  model,
  api,
  requested,
}: {
  userToken: string;
  model: string;
  api: APIFormat;
  requested: number;
}) {
  const user = users.get(userToken);
  if (!user) return false;
  if (user.type === "special") return true;

  const modelFamily = getModelFamilyForQuotaUsage(model, api);
  const { tokenCounts, tokenLimits } = user;
  const tokenLimit = tokenLimits[modelFamily];

  if (!tokenLimit) return true;

  const tokensConsumed = (tokenCounts[modelFamily] ?? 0) + requested;
  return tokensConsumed < tokenLimit;
}

export function refreshQuota(token: string) {
  const user = users.get(token);
  if (!user) return;
  const { tokenCounts, tokenLimits } = user;
  const quotas = Object.entries(config.tokenQuota) as [ModelFamily, number][];
  quotas
    // If a quota is not configured, don't touch any existing limits a user may
    // already have been assigned manually.
    .filter(([, quota]) => quota > 0)
    .forEach(
      ([model, quota]) =>
        (tokenLimits[model] = (tokenCounts[model] ?? 0) + quota)
    );
  usersToFlush.add(token);
}

export function resetUsage(token: string) {
  const user = users.get(token);
  if (!user) return;
  const { tokenCounts } = user;
  const counts = Object.entries(tokenCounts) as [ModelFamily, number][];
  counts.forEach(([model]) => (tokenCounts[model] = 0));
  usersToFlush.add(token);
}

/** Disables the given user, optionally providing a reason. */
export function disableUser(token: string, reason?: string) {
  const user = users.get(token);
  if (!user) return;
  user.disabledAt = Date.now();
  user.disabledReason = reason;
  usersToFlush.add(token);
}

export function getNextQuotaRefresh() {
  if (!quotaRefreshJob) return "never (manual refresh only)";
  return quotaRefreshJob.nextInvocation().getTime();
}

/**
 * Cleans up expired temporary tokens by disabling tokens past their access
 * expiry date and permanently deleting tokens three days after their access
 * expiry date.
 */
function cleanupExpiredTokens() {
  const now = Date.now();
  let disabled = 0;
  let deleted = 0;
  for (const user of users.values()) {
    if (user.type !== "temporary") continue;
    if (user.expiresAt && user.expiresAt < now && !user.disabledAt) {
      disableUser(user.token, "Temporary token expired.");
      disabled++;
    }
    if (user.disabledAt && user.disabledAt + 72 * 60 * 60 * 1000 < now) {
      users.delete(user.token);
      usersToFlush.add(user.token);
      deleted++;
    }
  }
  log.trace({ disabled, deleted }, "Expired tokens cleaned up.");
}

function refreshAllQuotas() {
  let count = 0;
  for (const user of users.values()) {
    if (user.type === "temporary") continue;
    refreshQuota(user.token);
    count++;
  }
  log.info(
    { refreshed: count, nextRefresh: quotaRefreshJob!.nextInvocation() },
    "Token quotas refreshed."
  );
}

// TODO: Firebase persistence is pretend right now and just polls the in-memory
// store to sync it with Firebase when it changes. Will refactor to abstract
// persistence layer later so we can support multiple stores.
let firebaseTimeout: NodeJS.Timeout | undefined;

async function initFirebase() {
  log.info("Connecting to Firebase...");
  const app = getFirebaseApp();
  const db = admin.database(app);
  const usersRef = db.ref("users");
  const snapshot = await usersRef.once("value");
  const users: Record<string, User> | null = snapshot.val();
  firebaseTimeout = setInterval(flushUsers, 20 * 1000);
  if (!users) {
    log.info("No users found in Firebase.");
    return;
  }
  for (const token in users) {
    upsertUser(users[token]);
  }
  usersToFlush.clear();
  const numUsers = Object.keys(users).length;
  log.info({ users: numUsers }, "Loaded users from Firebase");
}

async function flushUsers() {
  const app = getFirebaseApp();
  const db = admin.database(app);
  const usersRef = db.ref("users");
  const updates: Record<string, User> = {};
  const deletions = [];

  for (const token of usersToFlush) {
    const user = users.get(token);
    if (!user) {
      deletions.push(token);
      continue;
    }
    updates[token] = user;
  }

  usersToFlush.clear();

  const numUpdates = Object.keys(updates).length + deletions.length;
  if (numUpdates === 0) {
    return;
  }

  await usersRef.update(updates);
  await Promise.all(deletions.map((token) => usersRef.child(token).remove()));
  log.info(
    { users: Object.keys(updates).length, deletions: deletions.length },
    "Flushed changes to Firebase"
  );
}

function getModelFamilyForQuotaUsage(
  model: string,
  api: APIFormat
): ModelFamily {
  // "azure" here is added to model names by the Azure key provider to
  // differentiate between Azure and OpenAI variants of the same model.
  if (model.includes("azure")) return getAzureOpenAIModelFamily(model);
  if (model.includes("anthropic.")) return getAwsBedrockModelFamily(model);

  switch (api) {
    case "openai":
    case "openai-text":
    case "openai-image":
      return getOpenAIModelFamily(model);
    case "anthropic-chat":
    case "anthropic-text":
      return getClaudeModelFamily(model);
    case "google-ai":
      return getGoogleAIModelFamily(model);
    case "mistral-ai":
      return getMistralAIModelFamily(model);
    default:
      assertNever(api);
  }
}

function getRefreshCrontab() {
  switch (config.quotaRefreshPeriod!) {
    case "hourly":
      return "0 * * * *";
    case "daily":
      return "0 0 * * *";
    default:
      return config.quotaRefreshPeriod ?? "0 0 * * *";
  }
}


code:
import { Query } from "express-serve-static-core";
import sanitize from "sanitize-html";
import { z } from "zod";

export function parseSort(sort: Query["sort"]) {
  if (!sort) return null;
  if (typeof sort === "string") return sort.split(",");
  if (Array.isArray(sort)) return sort.splice(3) as string[];
  return null;
}

export function sortBy(fields: string[], asc = true) {
  return (a: any, b: any) => {
    for (const field of fields) {
      if (a[field] !== b[field]) {
        // always sort nulls to the end
        if (a[field] == null) return 1;
        if (b[field] == null) return -1;

        const valA = Array.isArray(a[field]) ? a[field].length : a[field];
        const valB = Array.isArray(b[field]) ? b[field].length : b[field];

        const result = valA < valB ? -1 : 1;
        return asc ? result : -result;
      }
    }
    return 0;
  };
}

export function paginate(set: unknown[], page: number, pageSize: number = 20) {
  const p = Math.max(1, Math.min(page, Math.ceil(set.length / pageSize)));
  return {
    page: p,
    items: set.slice((p - 1) * pageSize, p * pageSize),
    pageSize,
    pageCount: Math.ceil(set.length / pageSize),
    totalCount: set.length,
    nextPage: p * pageSize < set.length ? p + 1 : null,
    prevPage: p > 1 ? p - 1 : null,
  };
}

export function sanitizeAndTrim(
  input?: string | null,
  options: sanitize.IOptions = {
    allowedTags: [],
    allowedAttributes: {},
  }
) {
  return sanitize((input ?? "").trim(), options);
}

// https://github.com/colinhacks/zod/discussions/2050#discussioncomment-5018870
export function makeOptionalPropsNullable<Schema extends z.AnyZodObject>(
  schema: Schema
) {
  const entries = Object.entries(schema.shape) as [
    keyof Schema["shape"],
    z.ZodTypeAny
  ][];
  const newProps = entries.reduce(
    (acc, [key, value]) => {
      acc[key] =
        value instanceof z.ZodOptional ? value.unwrap().nullable() : value;
      return acc;
    },
    {} as {
      [key in keyof Schema["shape"]]: Schema["shape"][key] extends z.ZodOptional<
        infer T
      >
        ? z.ZodNullable<T>
        : Schema["shape"][key];
    }
  );
  return z.object(newProps);
}

export function redactIp(ip: string) {
  const ipv6 = ip.includes(":");
  return ipv6 ? "redacted:ipv6" : ip.replace(/\.\d+\.\d+$/, ".xxx.xxx");
}

export function assertNever(x: never): never {
  throw new Error(`Called assertNever with argument ${x}.`);
}


code:
<% if (flashData) {
  let flashStyle = { title: "", style: "" };
  switch (flashData.type) {
    case "success":
      flashStyle.title = "✅ Success:";
      flashStyle.style = "color: green; background-color: #ddffee; padding: 1em";
      break;
    case "error":
      flashStyle.title = "⚠ Error:";
      flashStyle.style = "color: red; background-color: #eedddd; padding: 1em";
      break;
    case "warning":
      flashStyle.title = "⚠ Alert:";
      flashStyle.style = "color: darkorange; background-color: #ffeecc; padding: 1em";
      break;
    case "info":
      flashStyle.title = "ℹ Notice:";
      flashStyle.style = "color: blue; background-color: #ddeeff; padding: 1em";
      break;
  }
%>
<p style="<%= flashStyle.style %>">
  <%= flashStyle.title %> <%= flashData.message %>

<% } %>


code:
<!DOCTYPE html>
<html>
  <head>
    <meta charset="utf-8">
    <meta name="csrf-token" content="<%= csrfToken %>">
    <title><%= title %></title>
    <style>
      a:hover {
        background-color: #e0e6f6;
      }

      a:visited:hover {
        background-color: #e7e0f6;
      }

      .pagination {
        list-style-type: none;
        padding: 0;
      }
      .pagination li {
        display: inline-block;
      }
      .pagination li a {
        display: block;
        padding: 0.5em 1em;
        text-decoration: none;
      }
      .pagination li.active a {
        background-color: #58739c;
        color: #fff;
      }

      table {
        border-collapse: collapse;
        border: 1px solid #ccc;
      }
      table.striped tr:nth-child(even) {
        background-color: #eaeaea
      }
      table td, table th {
        border: 1px solid #ccc;
        padding: 0.25em 0.5em;
      }
      th.active {
        background-color: #e0e6f6;
      }
      td.actions {
        padding: 0;
        width: 0;
        text-align: center;
      }
      td.actions a {
        text-decoration: none;
        background-color: transparent;
        padding: 0.5em;
        height: 100%;
        width: 100%;
      }
      td.actions:hover {
        background-color: #e0e6f6;
      }

      @media (max-width: 600px) {
        table {
          width: 100%;
        }
        table td, table th {
          display: block;
          width: 100%;
        }
      }
    </style>
  </head>
  <body style="font-family: sans-serif; background-color: #f0f0f0; padding: 1em;">
    <%- include("partials/shared_flash", { flashData: flash }) %>


code:
  Page Size
  <select id="pageSize" onchange="setPageSize(this.value)" style="margin-bottom: 1rem;">
    <option value="10" <% if (pageSize === 10) { %>selected<% } %>>10</option>
    <option value="20" <% if (pageSize === 20) { %>selected<% } %>>20</option>
    <option value="50" <% if (pageSize === 50) { %>selected<% } %>>50</option>
    <option value="100" <% if (pageSize === 100) { %>selected<% } %>>100</option>
    <option value="200" <% if (pageSize === 200) { %>selected<% } %>>200</option>
  </select>


code:
Next refresh: <time><%- nextQuotaRefresh %></time>

      Model Family
      Usage
      <% if (showTokenCosts) { %>
      Cost
      <% } %>
      Limit
      Remaining
      Refresh Amount

    <% Object.entries(quota).forEach(([key, limit]) => { %>

      <%- key %>
      <%- prettyTokens(user.tokenCounts[key]) %>
      <% if (showTokenCosts) { %>
      $<%- tokenCost(key, user.tokenCounts[key]).toFixed(2) %>
      <% } %>
      <% if (!user.tokenLimits[key]) { %>
      unlimited
      <% } else { %>
      <%- prettyTokens(user.tokenLimits[key]) %>
      <%- prettyTokens(user.tokenLimits[key] - user.tokenCounts[key]) %>
      <% } %>
      <% if (user.type === "temporary") { %>
      N/A
      <% } else { %>
      <%- prettyTokens(quota[key]) %>
      <% } %>

    <% }) %>


code:
Show all (<%- user.ip.length %>)

  <% user.ip.forEach((ip) => { %>

<%- shouldRedact ? redactIp(ip) : ip %>

  <% }) %>


code:
import cookieParser from "cookie-parser";
import expressSession from "express-session";
import MemoryStore from "memorystore";
import { config, COOKIE_SECRET } from "../config";

const ONE_WEEK = 1000 * 60 * 60 * 24 * 7;

const cookieParserMiddleware = cookieParser(COOKIE_SECRET);

const sessionMiddleware = expressSession({
  secret: COOKIE_SECRET,
  resave: false,
  saveUninitialized: false,
  store: new (MemoryStore(expressSession))({ checkPeriod: ONE_WEEK }),
  cookie: {
    sameSite: "strict",
    maxAge: ONE_WEEK,
    signed: true,
    secure: !config.useInsecureCookies,
  },
});

const withSession = [cookieParserMiddleware, sessionMiddleware];

export { withSession };


code:
import express, { Router } from "express";
import { injectCsrfToken, checkCsrfToken } from "../shared/inject-csrf";
import { selfServiceRouter } from "./web/self-service";
import { injectLocals } from "../shared/inject-locals";
import { withSession } from "../shared/with-session";

const userRouter = Router();

userRouter.use(
  express.json({ limit: "1mb" }),
  express.urlencoded({ extended: true, limit: "1mb" })
);
userRouter.use(withSession);
userRouter.use(injectCsrfToken, checkCsrfToken);
userRouter.use(injectLocals);

userRouter.use(selfServiceRouter);

userRouter.use(
  (
    err: Error,
    _req: express.Request,
    res: express.Response,
    _next: express.NextFunction
  ) => {
    const data: any = { message: err.message, stack: err.stack, status: 500 };
    res.status(500).render("user_error", { ...data, flash: null });
  }
);

export { userRouter };


code:
import { Router } from "express";
import { UserPartialSchema } from "../../shared/users/schema";
import * as userStore from "../../shared/users/user-store";
import { ForbiddenError, BadRequestError } from "../../shared/errors";
import { sanitizeAndTrim } from "../../shared/utils";
import { config } from "../../config";

const router = Router();

router.use((req, res, next) => {
  if (req.session.userToken) {
    res.locals.currentSelfServiceUser =
      userStore.getUser(req.session.userToken) || null;
  }
  next();
});

router.get("/", (_req, res) => {
  res.redirect("/");
});

router.get("/lookup", (_req, res) => {
  const ipLimit =
    (res.locals.currentSelfServiceUser?.maxIps ?? config.maxIpsPerUser) || 0;
  res.render("user_lookup", {
    user: res.locals.currentSelfServiceUser,
    ipLimit,
  });
});

router.post("/lookup", (req, res) => {
  const token = req.body.token;
  const user = userStore.getUser(token);
  req.log.info(
    { token: truncateToken(token), success: !!user },
    "User self-service lookup"
  );
  if (!user) {
    req.session.flash = { type: "error", message: "Invalid user token." };
    return res.redirect("/user/lookup");
  }
  req.session.userToken = user.token;
  return res.redirect("/user/lookup");
});

router.post("/edit-nickname", (req, res) => {
  const existing = res.locals.currentSelfServiceUser;

  if (!existing) {
    throw new ForbiddenError("Not logged in.");
  } else if (!config.allowNicknameChanges || existing.disabledAt) {
    throw new ForbiddenError("Nickname changes are not allowed.");
  } else if (!config.maxIpsAutoBan && !existing.ip.includes(req.ip)) {
    throw new ForbiddenError(
      "Nickname changes are only allowed from registered IPs."
    );
  }

  const schema = UserPartialSchema.pick({ nickname: true })
    .strict()
    .transform((v) => ({ nickname: sanitizeAndTrim(v.nickname) }));

  const result = schema.safeParse(req.body);
  if (!result.success) {
    throw new BadRequestError(result.error.message);
  }

  const newNickname = result.data.nickname || null;
  userStore.upsertUser({ token: existing.token, nickname: newNickname });
  req.session.flash = { type: "success", message: "Nickname updated." };
  return res.redirect("/user/lookup");
});

function truncateToken(token: string) {
  const sliceLength = Math.max(Math.floor(token.length / 8), 1);
  return `${token.slice(0, sliceLength)}...${token.slice(-sliceLength)}`;
}

export { router as selfServiceRouter };


code:
<%- include("partials/shared_header", { title: "Error" }) %>

⚠ Error <%= status %>: <%= message %>

<%= stack %>

      Go Back


code:
<!DOCTYPE html>

  <meta charset="UTF-8">
  <meta name="robots" content="noindex" />
  <title><%= title %></title>

    <%= pageHeader %>

Service Info

<%= JSON.stringify(serviceInfo, null, 2) %>


code:
<%- include("partials/shared_header", { title: "User Token Lookup" }) %>

User Token Lookup

Provide your user token to check your usage and quota information.

<form action="/user/lookup" method="post">
  <input type="hidden" name="_csrf" value="<%= csrfToken %>" />
  User Token
  <input type="password" name="token" value="<%= user?.token %>" />
  <input type="submit" value="Lookup" />
</form>
<% if (user) { %>

<% if (user.type === "temporary" && Boolean(user.disabledAt)) { %>
<%- include("partials/shared_flash", { flashData: {
      type: "info",
      message: "This temporary user token has expired and is no longer usable. These records will be deleted soon.",
    } }) %>
<% } else if (user.disabledAt) { %>
<%- include("partials/shared_flash", { flashData: {
      type: "warning",
      message: "This user token has been disabled." + (user.disabledReason ? ` Reason: ${user.disabledReason}` : ""),
    } }) %>
<% } %>

      User Token
       <%- "..." + user.token.slice(-5) %>

      Nickname
      <%- user.nickname ?? "none" %>

        ✏

      Type
      <%- user.type %>

      Prompts
      <%- user.promptCount %>

      Created At
      <%- user.createdAt %>

      Last Used At
      <%- user.lastUsedAt || "never" %>

      IPs<%- ipLimit ? ` (max ${ipLimit})` : "" %>
      <%- include("partials/shared_user_ip_list", { user, shouldRedact: true }) %>

    <% if (user.type === "temporary") { %>

      Expires At
      <%- user.expiresAt %>

    <% } %>

Quota Information

<%- include("partials/shared_quota-info", { quota, user }) %>

<form id="edit-nickname-form" style="display: none" action="/user/edit-nickname" method="post">
  <input type="hidden" name="_csrf" value="<%= csrfToken %>" />
  <input type="hidden" name="nickname" value="<%= user.nickname %>" />
</form>

<% } %> <%- include("partials/user_footer") %>


code:
*.bat	text eol=crlf
*.tw	text eol=lf
*.sh	text eol=lf
*.py	text eol=lf
*.txt	text eol=lf
compile	text eol=lf
compile-git	text eol=lf
sanityCheck	text eol=lf

path: .gitattributes


code:
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
*.swp
*.swo

# C extensions
*.so

# Distribution / packaging
.Python
env/
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
*.egg-info/
.installed.cfg
*.egg

# PyInstaller
#  Usually these files are written by a python script from a template
#  before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec

# Installer logs
pip-log.txt
pip-delete-this-directory.txt

# Unit test / coverage reports
htmlcov/
.tox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*,cover
.hypothesis/

# Translations
*.mo
*.pot

# Django stuff:
*.log
local_settings.py

# Flask stuff:
instance/
.webassets-cache

# Scrapy stuff:
.scrapy

# Sphinx documentation
docs/_build/

# PyBuilder
target/

# IPython Notebook
.ipynb_checkpoints

# pyenv
.python-version

# celery beat schedule file
celerybeat-schedule

# dotenv
.env

# virtualenv
venv/
ENV/

# Spyder project settings
.spyderproject

# Rope project settings
.ropeproject

# Start.tw
src/config/start.tw

path: .gitignore


code:
further development:
-specialized slave schools
-fortifications
-more levels for militia edict (further militarize society)
-conquering other arcologies?

Events:
-famous criminal escapes to the arcology, followed by another arcology police force

Bugs:
-sometimes troop counts breaks
-sometimes rebel numbers have fractionary parts


code:
# How to vector art

Please read the whole document before starting to work.
Some aspects are not obvious right away.

Note: This does not actually describe how to be an artist.

## TL;DR

    killall inkscape
    artTools/vector_layer_split.py artTools/vector_source.svg tw src/art/vector/layers/
    compile

## 1. Be an artist

Make changes to the vector_source.svg.
Inkscape was thoroughly tested.
Adobe Illustrator might work decently, too.

## 2. Respect the structure

While editing, keep the Layers in mind.

* In Inkscape, Layers are special "Inkscape groups".
* In Illustrator, Layers are groups with user-definded IDs.
* All Layers should have an ID that is globally unique
  (not just within their subtree).

* Please use anonymous groups only for the ease of editing. Remove all "helper" groups before finally saving the file.
* Anonymous groups can be used for continous scaling (of e.g. boobs).

* Every asset that should go into a separate file needs to be in a labelled group
  (even if that is a group with only one shape).
* There are some globally available styles defined as CSS classes (e.g. skin, hair).
  Use them if your asset should be changed upon display.
  Do not set the style directly in the object.

## 3. Normalise the document (before committing)

**THIS IS IMPORTANT**

The various editors out there (Inkscape, Illustrator) may use different variants of formatting the SVG XML.
Use

    python3 normalize_svg.py vector_source.svg

before committing to normalise the format so git will not freak out about the changed indentation.

If you use Inkscape, please use in Edit → Settings → Input/Output → SVG-Output → Path Data → Optimized. This is the default.
In case your Editor uses another path data style which cannot be changed, please contact the other artists and developers via the issue tracker to find a new common ground.

What it does:
* Formats the SVG XML according to Pythons lxml module, regardless of editor.
* Adobe Illustrator uses group IDs as layer labels.
  Inkscape however uses layer labels and a separate, auto-generated group ID.
  normalize_svg.py overwrites the group ID with the Inkscape layer label
  so they are synchronised with Inkscape layer labels.
* Inkscape copies the global style definition into the object *every time*
  the object is edited. If an object references a CSS class AND at the same time
  has a local style, the local style is removed
  so global dynamic styling is possible later on.

Note: Behaviour of Adobe Illustrator is untested.

## 4. Split the layers

For NoX original, deepurk extensions, execute

    python3 vector_layer_split.py vector_source_ndmain.svg tw ../src/art/vector/layers/

For faraen revamped art (based on NoX original)

	python3 vector_revamp_layer_split.py vector_revamp_source.svg tw ../src/art/vector_revamp/layers/

. This application reads all groups in `vector_source.svg`.
Each group is stored in a separate file in the target directory `/src/art/vector/layers/`.
The group ID sets the file name. Therefore, the group ID **must not** contain spaces or any other weird characters.

Also consider:
* A group with ID ending in _ is ignored. Use Layers ending in _ to keep overview.
* A group with ID starting with "g" (Inkscape) or "XMLID" (Illustrator) is also ignored.
* Existing files are overwritten **without enquiry**.
* The target directory is not emptied. If a file is no longer needed, you should remove it manually.
* This procedure removes global definitions. This means, SVG filters are currently not supported.

Available output formats are `svg` and `tw`.
`svg` output exists for debug reasons.
`tw` embeds the SVG data into Twine files, but removes the global style definitions so they can be set during display.

## 5. Edit the code

`/src/art/` contains Twine code which shows the assets in the story.
There are many helpful comments in `/src/art/artWidgets.tw`.
The code also generates the previously removed global style definitions on the fly and per display.

## 6. Compile the story

Use the provided compile script (Windows batch or shell) for compiling.

## 7. Advanced usage

You can define multiple CSS classes to one element, e.g. "skin torso". When procedurally generating CSS, you can then set a global "skin" style, and a different one for "torso".

You can put variable text into the image using single quotes, e.g. "'+_tattooText+'". The single quotes *break* the quoting in Twine so the content of the Twine variable `_tattooText` is shown upon display. You can even align text on paths.

An anonymous group can be used for dynamic transformation. However, finding the correct parameters for the affine transformation matrix can be cumbersome. Use any method you see fit. See `src/art/vector/Boob.tw` for one example of the result.


code:
Apron
AttractiveLingerie
AttractiveLingerieForAPregnantWoman
BallGown
Battledress
BodyOil
Bunny
Chains
ChattelHabit
Cheerleader
ClubslutNetting
ComfortableBodysuit
Conservative
Cutoffs
Cybersuit
FallenNunsHabit
FuckdollSuit
HalterTopDress
HaremGauze
Hijab
Huipil
Kimono
LatexCatsuit
Leotard
LongQipao
MaternityDress
MilitaryUniform
MiniDress
Monokini
NiceBusinessAttire
NiceMaid
NiceNurse
No
PenitentNunsHabit
RedArmyUniform
RestrictiveLatex
ScalemailBikini
SchoolgirlUniform
SchutzstaffelUniform
ShibariRopes
SlaveGown
Slutty
SluttyBusinessAttire
SluttyJewelry
SluttyMaid
SluttyNurse
SluttyQipao
SluttySchutzstaffelUniform
Spats
StretchPants
StringBikini
Succubus
Toga
UncomfortableStraps
Western


code:
#!/bin/bash

# This script reads all possible values of $slave.clothes as mentioned by the documentation.
# This script uses the actual implementation of the JS clothing2artSuffix function as defined in src/art/vector/Helper_Functions.tw
# This script outputs suffixes to be used in the SVG naming scheme
# This script is meant to be executed at the project root directory.
# This script depends on bash, grep, sed, paste and nodejs (so best executed on Linux, I guess)

# grep -Porh '(?<=\.clothes = )"[^"]+"' src/ | sort | uniq # searches sources for all clothes strings actually used in the game, even undocumented ones

(
  echo 'var window = {};'
  grep -v '^:' src/art/vector/Helper_Functions.tw
  echo -n 'Array('
  sed '/^clothes:/,/:/!d' "slave variables documentation - Pregmod.txt" | grep '"' | paste -sd,
  echo ').forEach(v => {console.log(window.clothing2artSuffix(v));});'
) | nodejs | sort


code:
#!/usr/bin/env python3

'''
Application for "normalizing" SVGs

These problems are addressed:

* Inkscape notoriously copies class styles into the object definitions.
https://bugs.launchpad.net/inkscape/+bug/167937

* Inkscape uses labels on layers. Layers are basically named groups.
  Inkscape does not sync the group id with the layer label.

Usage Example:
python3 inkscape_svg_fixup.py vector_source.svg
'''

import lxml.etree as etree
import sys

def fix(tree):
  # know namespaces
  ns = {
    'svg' : 'http://www.w3.org/2000/svg',
    'inkscape' : 'http://www.inkscape.org/namespaces/inkscape'
  }

  # find document global style definition
  # mangle it and interpret as python dictionary
  style_element = tree.find('./svg:style',namespaces=ns)
  style_definitions = style_element.text
  pythonic_style_definitions = '{'+style_definitions.\
  replace('\t.','"').replace('{','":"').replace('}','",').\
  replace('/*','#')+'}'
  styles = eval(pythonic_style_definitions)

  # go through all SVG elements
  for elem in tree.iter():
    if (elem.tag == etree.QName(ns['svg'], 'g')):
      # compare inkscape label with group element ID
      l = elem.get(etree.QName(ns['inkscape'], 'label'))
      if l:
        i = elem.get('id')
        if (i != l):
          print("Overwriting ID %s with Label %s..."%(i, l))
          elem.set('id', l)

    # clean styles (for easier manual merging)
    style_string = elem.get('style')
    if style_string:
      split_styles = style_string.strip('; ').split(';')
      styles_pairs = [s.strip('; ').split(':') for s in split_styles]
      filtered_pairs = [ (k,v) for k,v in styles_pairs if not (
        k.startswith('font-') or
        k.startswith('text-') or
        k.endswith('-spacing') or
        k in ["line-height", " direction", " writing", " baseline-shift", " white-space", " writing-mode"]
      )]
      split_styles = [':'.join(p) for p in filtered_pairs]
      style_string = ';'.join(sorted(split_styles))
      elem.attrib["style"] = style_string

    # remove all style attributes offending class styles
    s = elem.get('style')
    c = elem.get('class')
    if (c and s):
      s = s.lower()
      c = c.split(' ')[0] # regard main style only
      cs = ''
      if c in styles:
        cs = styles[c].strip('; ').lower()
      if (c not in styles):
        print("Object id %s references unknown style class %s."%(i,c))
      else:
        if (cs != s.strip('; ')):
          print("Style %s removed from object id %s differed from class %s style %s."%(s,i,c,cs))
        del elem.attrib["style"]

if __name__ == "__main__":
  input_file = sys.argv[1]
  tree = etree.parse(input_file)
  fix(tree)
  # store SVG into file (input file is overwritten)
  svg = etree.tostring(tree, pretty_print=True)
  with open(input_file, 'wb') as f:
    f.write('<?xml version="1.0" encoding="UTF-8" standalone="no"?>\n'.encode("utf-8"))
    f.write(svg)


code:
#!/usr/bin/env python3

'''
Application for procedural content adaption
*THIS IS VERY EXPERIMENTAL*
Contains a very poor man's implementation of spline mesh warping.

This application parses SVG path data of an outfit sample aligned to a body part.
The outfit is then replicated for other shapes of the same body part.

Example data is geared towards generating a strap outfit for boobs and torso
for all sizes of boobs and all shapes of torsos based on a single outfit for
boobs of size 2 and a hourglass torso respectively.

Limitations:
* handles paths only
* only svg.path Line and CubicBezier are tested
* only the aforementioned examples are tested

Usage:
python3 vector_clothing_replicator.py infile clothing bodypart destinationfile

Usage Example:
python3 vector_clothing_replicator.py vector_source.svg Straps Boob vector_destination.svg
python3 vector_clothing_replicator.py vector_source.svg Straps Torso vector_destination.svg
'''

from svg.path import parse_path
import copy
import lxml.etree as etree
import sys

REFERENCE_PATH_SAMPLES = 200
EMBED_REPLICATIONS = True # wether to embed all replications into the input file or output separate files

input_file = sys.argv[1]
clothing = sys.argv[2]
bodypart = sys.argv[3]
output_file_embed = sys.argv[4]

# TODO: make these configurable
output_file_pattern = '%s_%s_%s.svg' #bodypart, target_id, clothing

if ('Torso' == bodypart):
  xpath_shape = './svg:g[@id="Torso_"]/svg:g[@id="Torso_%s"]/svg:path[@class="skin torso"]/@d' # TODO: formulate more general, independent of style
  xpath_outfit_container = '//svg:g[@id="Torso_Outfit_%s_"]'%(clothing)
  xpath_outfit = '//svg:g[@id="Torso_Outfit_%s_%s"]'%(clothing,'%s')
  target_ids = "Unnatural,Hourglass,Normal".split(",")
  reference_id = "Hourglass"
else:
  raise RuntimeError("Please specify a bodypart for clothing to replicate.")

tree = etree.parse(input_file)
ns = {'svg' : 'http://www.w3.org/2000/svg'}

canvas = copy.deepcopy(tree)
for e in canvas.xpath('./svg:g',namespaces=ns)+canvas.xpath('./svg:path',namespaces=ns):
  # TODO: this should be "remove all objects, preserve document properties"
  e.getparent().remove(e)

def get_points(xpath_shape):
  '''
  This funciton extracts reference paths by the given xpath selector.
  Each path is used to sample a fixed number of points.
  '''
  paths_data = tree.xpath(xpath_shape,namespaces=ns)
  points = []
  path_length = None
  for path_data in paths_data:
    p = parse_path(path_data)
    points += [
      p.point(1.0/float(REFERENCE_PATH_SAMPLES)*i)
      for i in range(REFERENCE_PATH_SAMPLES)
    ]
  if (not points):
    raise RuntimeError(
      'No paths for reference points found by selector "%s".'%(xpath_shape)
    )
  return points

def point_movement(point, reference_points, target_points):
  '''
  For a given point, finds the nearest point in the reference path.
  Gives distance vector from the nearest reference point to the
  respective target reference point.
  '''
  distances = [abs(point-reference_point) for reference_point in reference_points]
  min_ref_dist_idx = min(enumerate(distances), key=lambda x:x[1])[0]
  movement = target_points[min_ref_dist_idx] - reference_points[min_ref_dist_idx]
  return movement

reference_points = get_points(xpath_shape%(reference_id))
container = tree.xpath(xpath_outfit_container,namespaces=ns)
if (len(container) != 1):
  raise RuntimeError('Outfit container selector "%s" does not yield exactly one layer.'%(xpath_outfit_container))
container = container[0]
outfit_source = container.xpath(xpath_outfit%(reference_id),namespaces=ns)
if (len(outfit_source) != 1):
  raise RuntimeError('Outfit source selector "%s" does not yield exactly one outfit layer in container selected by "%s".'%(xpath_outfit%(reference_id), xpath_outfit_container))
outfit_source = outfit_source[0]

for target_id in target_ids:
  print(
    'Generating variant "%s" of clothing "%s" for bodypart "%s"...'%
    (target_id, clothing, bodypart)
  )
  outfit = copy.deepcopy(outfit_source)
  paths = outfit.xpath('./svg:path',namespaces=ns)
  if target_id == reference_id:
    print("This is the source variant. Skipping...")
  else:
    layerid = outfit.get('id').replace('_%s'%(reference_id),'_%s'%(target_id))
    outfit.set('id', layerid)
    outfit.set(etree.QName('http://www.inkscape.org/namespaces/inkscape', 'label'), layerid) # for the Inkscape-users
    target_points = get_points(xpath_shape%(target_id))
    if (len(reference_points) != len(target_points)):
      raise RuntimeError(
        ('Different amounts of sampled points in reference "%s" and target "%s" paths. '+
        'Selector "%s" probably matches different number of paths in the two layers.')%
        (reference_id, target_id, xpath_shape)
      )
    for path in paths:
      path_data = path.get("d")
      p = parse_path(path_data)
      for segment in p:
        original_distance = abs(segment.end-segment.start)
        start_movement = point_movement(segment.start, reference_points, target_points)
        segment.start += start_movement
        end_movement = point_movement(segment.end, reference_points, target_points)
        segment.end += end_movement
        distance = abs(segment.end-segment.start)
        try:
          # enhance position of CubicBezier control points
          # amplification is relative to the distance gained by movement
          segment.control1 += start_movement
          segment.control1 += (segment.control1-segment.start)*(distance/original_distance-1.0)
          segment.control2 += end_movement
          segment.control2 += (segment.control2-segment.end)*(distance/original_distance-1.0)
        except AttributeError as ae:
          # segment is not a CubicBezier
          pass
      path.set("d", p.d())
    if EMBED_REPLICATIONS:
      container.append(outfit)
  if not EMBED_REPLICATIONS:
    container = copy.deepcopy(canvas).xpath('.',namespaces=ns)[0]
    container.append(outfit)

  if not EMBED_REPLICATIONS:
    svg = etree.tostring(container, pretty_print=True)
    with open((output_file_pattern%(bodypart, target_id, clothing)).lower(), 'wb') as f:
      f.write('<?xml version="1.0" encoding="UTF-8" standalone="no"?>\n'.encode("utf-8"))
      f.write(svg)

if EMBED_REPLICATIONS:
  svg = etree.tostring(tree, pretty_print=True)
  with open(output_file_embed, 'wb') as f:
    f.write('<?xml version="1.0" encoding="UTF-8" standalone="no"?>\n'.encode("utf-8"))
    f.write(svg)


code:
#!/usr/bin/env python3

'''
Application for splitting groups from one SVG file into separate files

Usage:
python3 vector_layer_split.py infile format outdir

Usage Example:
python3 vector_layer_split.py vector_source.svg tw ../src/art/vector/layers/
'''

import lxml.etree as etree
import sys
import os
import copy
import re
import normalize_svg

input_file = sys.argv[1]
output_format = sys.argv[2]
output_directory = sys.argv[3]
if not os.path.exists(output_directory):
    os.makedirs(output_directory)

ns = {
  'svg' : 'http://www.w3.org/2000/svg',
  'inkscape' : 'http://www.inkscape.org/namespaces/inkscape',
  'sodipodi':"http://sodipodi.sourceforge.net/DTD/sodipodi-0.dtd",
}

tree = etree.parse(input_file)
normalize_svg.fix(tree)

# prepare output template
template = copy.deepcopy(tree)
root = template.getroot()

# remove all svg root attributes except document size
for a in root.attrib:
  if (a != "viewBox"):
    del root.attrib[a]

# add placeholder for CSS class (needed for workaround for non HTML 5.1 compliant browser)
if output_format == 'tw':
  root.attrib["class"] = "'+_art_display_class+'"

# remove all content, including metadata
# for twine output, style definitions are removed, too
defs = None
for e in root:
  if (e.tag == etree.QName(ns['svg'], 'defs')):
    defs = e
  if (e.tag == etree.QName(ns['svg'], 'g') or
      e.tag == etree.QName(ns['svg'], 'metadata') or
      e.tag == etree.QName(ns['svg'], 'defs') or
      e.tag == etree.QName(ns['sodipodi'], 'namedview') or
      (output_format == 'tw' and e.tag == etree.QName(ns['svg'], 'style'))
      ):
    root.remove(e)
# template preparation finished

# prepare regex for later use
regex_xmlns = re.compile(' xmlns[^ ]+',)
regex_space = re.compile('[>][ ]+[<]',)

# find all groups
layers = tree.xpath('//svg:g',namespaces=ns)
for layer in layers:
  i = layer.get('id')
  if ( # disregard non-content groups
    i.endswith("_") or # manually suppressed with underscore
    i.startswith("XMLID") or # Illustrator generated group
    i.startswith("g") # Inkscape generated group
  ):
    continue
  # create new canvas
  output = copy.deepcopy(template)
  # copy all shapes into template
  canvas = output.getroot()
  for e in layer:
    canvas.append(e)
  # represent template as SVG (binary string)
  svg = etree.tostring(output, pretty_print=False)
  # poor man's conditional defs insertion
  # TODO: extract only referenced defs (filters, gradients, ...)
  # TODO: detect necessity by traversing the elements. do not stupidly search in the string representation
  if ("filter:" in svg.decode('utf-8')):
    # it seems there is a filter referenced in the generated SVG, re-insert defs from main document
    canvas.insert(0,defs)
    # re-generate output
    svg = etree.tostring(output, pretty_print=False)

  if (output_format == 'tw'):
    # remove unnecessary attributes
    # TODO: never generate unnecessary attributes in the first place
    svg = svg.decode('utf-8')
    svg = regex_xmlns.sub('',svg)
    svg = svg.replace(' inkscape:connector-curvature="0"','') # this just saves space
    svg = svg.replace('\n','').replace('\r','') # print cannot be multi-line
    svg = regex_space.sub('><',svg) # remove indentaion
    svg = svg.replace('svg:','') # svg namespace was removed
    svg = svg.replace('<g ','<g transform="\'+_art_transform+\'"') # internal groups are used for scaling
    svg = svg.encode('utf-8')

  # save SVG string to file
  i = layer.get('id')
  output_path = os.path.join(output_directory,i+'.'+output_format)
  with open(output_path, 'wb') as f:
    if (output_format == 'svg'):
      # Header for normal SVG (XML)
      f.write('<?xml version="1.0" encoding="UTF-8" standalone="no"?>\n'.encode("utf-8"))
      f.write(svg)
    elif (output_format == 'tw'):
      # Header for SVG in Twine file (SugarCube print statement)
      f.write((':: Art_Vector_%s [nobr]\n\n'%(i)).encode("utf-8"))
      f.write("<<print '<html>".encode("utf-8"))
      f.write(svg)
      f.write("</html>' >>".encode("utf-8"))


code:
#!/usr/bin/env python3

'''
Application for splitting groups from one SVG file into separate files

Usage:
python3 vector_layer_split.py infile format outdir

Usage Example:
python3 vector_layer_split.py vector_source.svg tw ../src/art/vector/layers/
'''

import lxml.etree as etree
from lxml.etree import XMLParser, parse

import sys
import os
import copy
import re
import normalize_svg

input_file = sys.argv[1]
output_format = sys.argv[2]
output_directory = sys.argv[3]

if not os.path.exists(output_directory):
    os.makedirs(output_directory)

ns = {
    'svg': 'http://www.w3.org/2000/svg',
    'inkscape': 'http://www.inkscape.org/namespaces/inkscape',
    'sodipodi': "http://sodipodi.sourceforge.net/DTD/sodipodi-0.dtd",
}

p = XMLParser(huge_tree=True)
tree = parse(input_file, parser=p)
#tree = etree.parse(input_file)
normalize_svg.fix(tree)

# prepare output template
template = copy.deepcopy(tree)
root = template.getroot()

# remove all svg root attributes except document size
for a in root.attrib:
    if (a != "viewBox"):
        del root.attrib[a]

# add placeholder for CSS class (needed for workaround for non HTML 5.1 compliant browser)
if output_format == 'tw':
    root.attrib["class"] = "'+_art_display_class+'"

# remove all content, including metadata
# for twine output, style definitions are removed, too
defs = None
for e in root:
    if (e.tag == etree.QName(ns['svg'], 'defs')):
        defs = e
    if (e.tag == etree.QName(ns['svg'], 'g') or
            e.tag == etree.QName(ns['svg'], 'metadata') or
            e.tag == etree.QName(ns['svg'], 'defs') or
            e.tag == etree.QName(ns['sodipodi'], 'namedview') or
            (output_format == 'tw' and e.tag == etree.QName(ns['svg'], 'style'))
    ):
        root.remove(e)
# template preparation finished

# prepare regex for later use
regex_xmlns = re.compile(' xmlns[^ ]+', )
regex_space = re.compile('[>][ ]+[<]', )
regex_transformValue = re.compile('(?<=transformVariableName=")[^"]*', )
regex_transformVar = re.compile('transformVariableName="[^"]*"', )
regex_transformAttr = re.compile('transform="[^"]*"', )
# find all groups
layers = tree.xpath('//svg:g', namespaces=ns)
for layer in layers:
    i = layer.get('id')
    if (  # disregard non-content groups
            i.endswith("_") or  # manually suppressed with underscore
            i.startswith("XMLID") or  # Illustrator generated group
            i.startswith("g")  # Inkscape generated group
    ):
        continue
    # create new canvas
    output = copy.deepcopy(template)
    # copy all shapes into template
    canvas = output.getroot()
    for e in layer:
        canvas.append(e)
    # represent template as SVG (binary string)
    svg = etree.tostring(output, pretty_print=False)
    # poor man's conditional defs insertion
    # TODO: extract only referenced defs (filters, gradients, ...)
    # TODO: detect necessity by traversing the elements. do not stupidly search in the string representation
    if ("filter:" in svg.decode('utf-8')):
        # it seems there is a filter referenced in the generated SVG, re-insert defs from main document
        canvas.insert(0, defs)
        # re-generate output
        svg = etree.tostring(output, pretty_print=False)
    elif ("clip-path=" in svg.decode('utf-8')):
        # it seems there is a clip path referenced in the generated SVG, re-insert defs from main document
        canvas.insert(0, defs)
        # re-generate output
        svg = etree.tostring(output, pretty_print=False)

    if (output_format == 'tw'):
        # remove unnecessary attributes
        # TODO: never generate unnecessary attributes in the first place
        svg = svg.decode('utf-8')
        svg = regex_xmlns.sub('', svg)
        svg = svg.replace(' inkscape:connector-curvature="0"', '')  # this just saves space
        svg = svg.replace('\n', '').replace('\r', '')  # print cannot be multi-line
        svg = regex_space.sub('><', svg)  # remove indentaion
        svg = svg.replace('svg:', '')  # svg namespace was removed

        transformGroups = regex_transformVar.findall(svg)

        if (len(transformGroups) > 0):
            svg = regex_transformAttr.sub('', svg)
            for transformGroup in transformGroups:
                transformValue = regex_transformValue.search(transformGroup)
                if (transformValue is not None):
                    svg = svg.replace(transformGroup, ' transform="\'+' + transformValue.group() + '+\'" ')  # internal groups are used for scaling

        svg = svg.encode('utf-8')

    # save SVG string to file
    i = layer.get('id')
    output_path = os.path.join(output_directory, i + '.' + output_format)
    with open(output_path, 'wb') as f:
        if (output_format == 'svg'):
            # Header for normal SVG (XML)
            f.write('<?xml version="1.0" encoding="UTF-8" standalone="no"?>\n'.encode("utf-8"))
            f.write(svg)
        elif (output_format == 'tw'):
            # Header for SVG in Twine file (SugarCube print statement)
            f.write((':: Art_Vector_Revamp_%s [nobr]\n\n' % (i)).encode("utf-8"))
            f.write("<<print '<html>".encode("utf-8"))
            f.write(svg)
            f.write("</html>' >>".encode("utf-8"))


code:
# Ignore everything in this directory
*
# Except the following:
!.gitignore


code:
#!/bin/bash

while [[ "$1" ]]
do
	case $1 in
	--insane)
		insane="true"
	;;
	*)
		echo "Unknown argument $1"
		exit 1
	esac
	shift
done

# Find and insert current commit
COMMIT=$(git rev-parse --short HEAD)
sed -Ei "s/build .releaseID/\0 commit $COMMIT/" src/gui/mainMenu/AlphaDisclaimer.tw

if [[ ! "$insane" ]]
then
	# Run sanity check.
	./sanityCheck
fi

ARCH="$(uname -m)"
if [ "$ARCH" = "x86_64" ]
then
	echo "x64 arch"
	./devTools/tweeGo/tweego_nix64 -o bin/FC_pregmod.html src/ || build_failed="true"
elif echo "$ARCH" | grep -Ee '86$' > /dev/null
then
	echo "x86 arch"
	./devTools/tweeGo/tweego_nix86 -o bin/FC_pregmod.html src/ || build_failed="true"
elif echo "$ARCH" | grep -Ee '^arm' > /dev/null
then
	echo "arm arch"
	# tweego doesn't provide arm binaries so you have to build it yourself
	export TWEEGO_PATH=devTools/tweeGo/storyFormats
	tweego -o bin/FC_pregmod.html src/ || build_failed="true"
else
	exit 2
fi

# Revert AlphaDisclaimer for next compilation
git checkout -- src/gui/mainMenu/AlphaDisclaimer.tw

if [ "$build_failed" = "true" ]
then
	exit 1
fi

#Make the output prettier, replacing \t with a tab and \n with a newline
sed -i -e '/^.*<div id="store-area".*$/s/\\t/\t/g' -e '/^.*<div id="store-area".*$/s/\\n/\n/g'  bin/FC_pregmod.html

path: compile


code:
#!/bin/bash

# Run sanity check.
./sanityCheck

HASH="$(git rev-list -n 1 --abbrev-commit HEAD)"
ARCH="$(uname -m)"
if [ "$ARCH" = "x86_64" ]
then
	echo "x64 arch"
	./devTools/tweeGo/tweego_nix64 -o "bin/FC_pregmod_$HASH.html" src/
elif echo "$ARCH" | grep -Ee '86$' > /dev/null
then
	echo "x86 arch"
	./devTools/tweeGo/tweego_nix86 -o "bin/FC_pregmod_$HASH.html" src/
elif echo "$ARCH" | grep -Ee '^arm' > /dev/null
then
	echo "arm arch"
	# tweego doesn't provide arm binaries so you have to build it yourself
	export TWEEGO_PATH=devTools/tweeGo/storyFormats
	tweego -o "bin/FC_pregmod_$HASH.html" src/
else
	exit 2
fi

#Make the output prettier, replacing \t with a tab and \n with a newline
sed -i -e '/^<div id="store-area".*$/s/\\t/\t/g' -e '/^<div id="store-area".*$/s/\\n/\n/g'  "bin/FC_pregmod_$HASH.html"

echo "FC_pregmod_$HASH.html" compilation finished.

path: compile-git


code:
@echo off
:: Free Cities Basic Compiler - Windows

:: Set working directory
pushd %~dp0

:: Run the appropriate compiler for the user's CPU architecture.
if %PROCESSOR_ARCHITECTURE% == AMD64 (
    CALL "%~dp0devTools\tweeGo\tweego_win64.exe" -o "%~dp0bin/FC_pregmod.html" "%~dp0src"
) else (
    CALL "%~dp0devTools\tweeGo\tweego_win86.exe" -o "%~dp0bin/FC_pregmod.html" "%~dp0src"
)

popd
ECHO Done


code:
@echo off
:: Free Cities Basic Compiler - Windows

:: Set working directory
pushd %~dp0

:: See if we can find a git installation
setlocal enabledelayedexpansion

for %%k in (HKCU HKLM) do (
    for %%w in (\ \Wow6432Node\) do (
        for /f "skip=2 delims=: tokens=1*" %%a in ('reg query "%%k\SOFTWARE%%wMicrosoft\Windows\CurrentVersion\Uninstall\Git_is1" /v InstallLocation 2^> nul') do (
            for /f "tokens=3" %%z in ("%%a") do (
                set GIT=%%z:%%b
				set GITFOUND=yes
                goto FOUND
            )
        )
    )
)
:FOUND
if %GITFOUND% == yes (
	set "PATH=%GIT%bin;%PATH%"
	bash --login -c ./sanityCheck
)

:: Compile the game
call "%~dp0compile.bat"

if %GITFOUND% == yes (
	:: Make the output prettier, replacing \t with a tab and \n with a newline
	bash -c "sed -i -e '/^.*<div id=\"store-area\".*$/s/\\\t/\t/g' -e '/^.*<div id=\"store-area\".*$/s/\\\n/\n/g' bin/FC_pregmod.html"

	:: Revert ./src/init/storyInit.tw for next compilation
	git checkout -- ./src/init/storyInit.tw
)

popd
PAUSE


code:
@echo off
:: Free Cities Basic Compiler - Windows

:: Set working directory
pushd %~dp0

:: Compile the game
call "%~dp0compile.bat"

popd
PAUSE