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(0, 50000),
      entry.promptFlattened.slice(0, 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.info("Creating new index sheet.");
    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: AllieB/pacman

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(e, "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 { 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 "gpt4-32k":
      cost = 0.00006;
      break;
    case "gpt4":
      cost = 0.00003;
      break;
    case "turbo":
      cost = 0.0000015;
      break;
    case "claude":
      cost = 0.00001102;
      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(2) + "b";
  }
}


code:
import { getTokenizer } from "@anthropic-ai/tokenizer";
import { Tiktoken } from "tiktoken/lite";

let encoder: Tiktoken;

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();
  return true;
}

export function getTokenCount(prompt: string, _model: string) {
  // Don't try tokenizing if the prompt is massive to prevent DoS.
  // 500k characters should be sufficient for all supported models.
  if (prompt.length > 500000) {
    return {
      tokenizer: "length fallback",
      token_count: 100000,
    };
  }

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


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


code:
import { Tiktoken } from "tiktoken/lite";
import cl100k_base from "tiktoken/encoders/cl100k_base.json";

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 function getTokenCount(
  prompt: string | OpenAIPromptMessage[],
  model: string
) {
  if (typeof prompt === "string") {
    return getTextTokenCount(prompt);
  }

  const gpt4 = model.startsWith("gpt-4");

  const tokensPerMessage = gpt4 ? 3 : 4;
  const tokensPerName = gpt4 ? 1 : -1; // turbo omits role if name is present

  let numTokens = 0;

  for (const message of prompt) {
    numTokens += tokensPerMessage;
    for (const key of Object.keys(message)) {
      {
        const value = message[key as keyof OpenAIPromptMessage];
        if (!value || typeof value !== "string") continue;
        // Break if we get a huge message or exceed the token limit to prevent
        // DoS.
        // 100k tokens allows for future 100k GPT-4 models and 500k characters
        // is just a sanity check
        if (value.length > 500000 || numTokens > 100000) {
          numTokens = 100000;
          return {
            tokenizer: "tiktoken (prompt length limit exceeded)",
            token_count: numTokens,
          };
        }

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

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

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

export type OpenAIPromptMessage = {
  name?: string;
  content: string;
  role: string;
};


code:
import { Request } from "express";
import { config } from "../../config";
import { assertNever } from "../utils";
import {
  init as initClaude,
  getTokenCount as getClaudeTokenCount,
} from "./claude";
import {
  init as initOpenAi,
  getTokenCount as getOpenAITokenCount,
  OpenAIPromptMessage,
} from "./openai";
import { APIFormat } from "../key-management";

export async function init() {
  if (config.anthropicKey) {
    initClaude();
  }
  if (config.openaiKey || config.googlePalmKey) {
    initOpenAi();
  }
}

/** 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 } & (
  | { prompt: OpenAIPromptMessage[]; completion?: never; service: "openai" }
  | {
      prompt: string;
      completion?: never;
      service: "openai-text" | "anthropic" | "google-palm";
    }
  | { prompt?: never; completion: string; service: APIFormat }
);

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":
      return {
        ...getClaudeTokenCount(prompt ?? completion, req.body.model),
        tokenization_duration_ms: getElapsedMs(time),
      };
    case "openai":
    case "openai-text":
      return {
        ...getOpenAITokenCount(prompt ?? completion, req.body.model),
        tokenization_duration_ms: getElapsedMs(time),
      };
    case "google-palm":
      // TODO: Can't find a tokenization library for PaLM. There is an API
      // endpoint for it but it adds significant latency to the request.
      return {
        ...getOpenAITokenCount(prompt ?? completion, req.body.model),
        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),
  claude: z.number().optional().default(0),
  bison: 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 { ModelFamily } from "../models";
import { logger } from "../../logger";
import { User, UserTokenCounts, UserUpdate } from "./schema";

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

const INITIAL_TOKENS: Required<UserTokenCounts> = {
  turbo: 0,
  gpt4: 0,
  "gpt4-32k": 0,
  claude: 0,
  bison: 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) {
    updates.tokenCounts["gpt4-32k"] ??= 0;
    updates.tokenCounts["bison"] ??= 0;
  }
  if (updates.tokenLimits) {
    updates.tokenLimits["gpt4-32k"] ??= 0;
    updates.tokenLimits["bison"] ??= 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,
  consumption: number
) {
  const user = users.get(token);
  if (!user) return;
  const modelFamily = getModelFamilyForQuotaUsage(model);
  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) {
  const user = users.get(token);
  if (!user || user.disabledAt) return;
  if (!user.ip.includes(ip)) user.ip.push(ip);

  const configIpLimit = user.maxIps ?? config.maxIpsPerUser;
  const ipLimit =
    user.type === "special" || !configIpLimit ? Infinity : configIpLimit;
  if (user.ip.length > ipLimit) {
    disableUser(token, "IP address limit exceeded.");
    return;
  }

  user.lastUsedAt = Date.now();
  usersToFlush.add(token);
  return user;
}

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

  const modelFamily = getModelFamilyForQuotaUsage(model);
  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.debug({ 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"
  );
}

// TODO: use key-management/models.ts for family mapping
function getModelFamilyForQuotaUsage(model: string): ModelFamily {
  if (model.includes("32k")) {
    return "gpt4-32k";
  }
  if (model.startsWith("gpt-4")) {
    return "gpt4";
  }
  if (model.startsWith("gpt-3.5")) {
    return "turbo";
  }
  if (model.includes("bison")) {
    return "bison";
  }
  if (model.includes("claude")) {
    return "claude";
  }
  throw new Error(`Unknown quota model family for model ${model}`);
}

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 { 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 },
});

const withSession = [cookieParserMiddleware, sessionMiddleware];

export { withSession };


code:
import { Express } from "express-serve-static-core";
import { APIFormat, Key } from "../shared/key-management/index";
import { User } from "../shared/users/user-store";

declare global {
  namespace Express {
    interface Request {
      key?: Key;
      /** Denotes the format of the user's submitted request. */
      inboundApi: APIFormat;
      /** Denotes the format of the request being proxied to the API. */
      outboundApi: APIFormat;
      /** If the request comes from a RisuAI.xyz user, this is their token. */
      risuToken?: string;
      user?: User;
      isStreaming?: boolean;
      startTime: number;
      retryCount: number;
      queueOutTime?: number;
      onAborted?: () => void;
      proceed: () => void;
      heartbeatInterval?: NodeJS.Timeout;
      promptTokens?: number;
      outputTokens?: number;
      // TODO: remove later
      debug: Record<string, any>;
    }
  }
}

declare module "express-session" {
  interface SessionData {
    adminToken?: string;
    userToken?: string;
    csrf?: string;
    flash?: { type: string; message: string };
  }
}


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, UserInputError } 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);
  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.");
  }

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

  const result = schema.safeParse(req.body);
  if (!result.success) {
    throw new UserInputError(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");
});

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:
# Submissive Partners

## Overview

This mod focuses on creating engaging
and rewarding interactions with the player’s sexual and romantic
partners centered around BDSM themes, primarily domination/submission
and control. My aim is to introduce mechanics that complement vanilla
gameplay and will mesh well with other mods, while still being
relatively lightweight. This will greatly expand how much you as both a
player and a ruler can get out of your love life, while avoiding pulling
 your attention away from regular gameplay.

Submissive Partners should appeal to
players with interests across the BDSM spectrum. Things will start off
light, and if that's as far as you want to go, that's fine! Care will be
 taken to tailor content to what you've indicated that you're interested
 in, and you can feel safe that you won't suddenly get hit with
something intense that you weren't prepared for. If you *are* looking
 for more hardcore kinks, you'll be able to push your relationships in
that direction, but it will be up to you to do so.

## Highlights

- Train spouses, concubines, and lovers
  in obedience and submissiveness, molding their personalities to suit
  your tastes and turning them into loyal servants for life

- Control your partner’s clothing privileges, orgasm privileges, and fertility

- As your partners progress in their
  training, unlock new rules you can set for them, new types of
  interactions, and more specialized ways to train them

- Earn various mechanical benefits from
  having well-trained partners (primarily fertility, prestige, stress
  loss, and vassal opinion) without breaking the game

## Requirements

- [Carnalitas](https://www.loverslab.com/files/file/14207-carnalitas-unified-sex-mod-framework-for-ck3/): Make sure it is active & above Submissive Partners in your load order

## Installating/Updating

If your CK3 mod folder doesn't exist yet (on Windows, Documents/Paradox
Interactive/Crusader Kings III/mod), you need to create it. To install,
unzip into your CK3 mod folder, add this mod to your playset in the CK3
launcher, and enable it.

If updating, I highly recommend you **delete the old mod files** and reinstall every time, in case I renamed something or deleted old files. Versions 0.3+ are **incompatible** with save games from older versions.

## Compatibility

No known compatibility issues, other than with mods using the same trait indices as mine (5600-5699). Be sure to check the compatibility section for Carnalitas, though.

## Contributing

If you are interested in contributing, get in touch with me via either of the links above. Currently this repo isn't really set up to handle contributions from multiple authors (as I don't expect anyone else to want to write code for it, really), but that can be changed if there's a reason to do so.

## Credits

- Caffe: Lead Developer & Author
- erri120: Creator of localization_helper
- Irivial: French translation
- Triskelia: Icon artwork
- gallowsuk: bugfix

Submissive Partners takes heavy inspiration from [Immersive Slave Training and Kidnapping](https://www.loverslab.com/topic/136578-mod-immersive-slave-training-and-kidnapping/) and, by extension, [Kings of Slaves and Vices](https://www.loverslab.com/files/file/2574-ck2-kings-of-slaves-and-vices/) and [KSV Remastered](https://www.loverslab.com/topic/146590-mod-ck2-kings-of-slaves-and-vices-remastered/), so a big thank you to the creators of all 3 of those as well, even if they didn't directly contribute!

## Links

[Discord](https://discord.gg/FvqTR7u)

[Loverslab thread](https://www.loverslab.com/topic/151976-mod-ck3-submissive-partners/)

[GitHub (Original CaffeGit repository)]([GitHub - CaffeGit/Submissive-Partners](https://github.com/CaffeGit/Submissive-Partners))

[GitGud repository](https://gitgud.io/crusader-kings-3/Submissive-Partners)

repo_name: crusader/kings

license: mit


code:
version="1.2.2"
tags={
	"Utilities"
}
name="Submissive Partners"
supported_version="1.9.*"
path="mod/Submissive Partners"

path: Submissive Partners.mod


code:
#######
### RULES/SIMPLE INTERACTIONS
#######

# Set Clothing Rules
sub_clothing_rules_interaction = {
	category = interaction_category_friendly
	common_interaction = yes
	interface_priority = 96

	desc = sub_clothing_rules_interaction_desc

	auto_accept = yes

	is_shown = {
		scope:recipient = {
			sub_show_sub_interactions_trigger = { DOMINANT = scope:actor }
		}
	}

	is_valid_showing_failures_only = {
		scope:recipient = { has_trait = sub_subservient }
	}

	on_accept = {
		scope:actor = {
			trigger_event = {
				on_action = sub_set_clothing_rules
			}
		}
	}
}

# Set Chastity Rules
sub_chastity_rules_interaction = {
	category = interaction_category_friendly
	common_interaction = yes
	interface_priority = 95

	desc = sub_chastity_rules_interaction_desc

	auto_accept = yes

	is_shown = {
		scope:recipient = {
			sub_show_sub_interactions_trigger = { DOMINANT = scope:actor }
		}
	}

	is_valid_showing_failures_only = {
		scope:recipient = { has_trait = sub_subservient }
	}

	on_accept = {
		scope:actor = {
			trigger_event = {
				on_action = sub_set_chastity_rules
			}
		}
	}
}

# Posing
sub_show_pose_interaction = {
	category = interaction_category_friendly
	common_interaction = no
	interface_priority = 90

	desc = sub_show_pose_interaction_desc

	auto_accept = yes

	is_shown = {
		scope:recipient = {
			sub_show_sub_interactions_trigger = { DOMINANT = scope:actor }
		}
	}

	is_valid_showing_failures_only = {
		scope:recipient = { has_trait = sub_subservient }
	}

	on_accept = {
		scope:actor = {
			trigger_event = {
				on_action = sub_ask_to_pose
			}
		}
	}
}

# Remove submissive relation
sub_remove_submissive_interaction = {
	category = interaction_category_hostile

	auto_accept = yes
	ignores_pending_interaction_block = yes

	desc = sub_remove_submissive_interaction_desc

	is_shown = {
		scope:actor = {
			has_relation_submissive = scope:recipient
		}
	}

	on_accept = {
		scope:actor = {
			send_interface_message = {
				type = event_generic_neutral
				title = sub_remove_submissive_interaction_notification
				right_icon = scope:recipient
				sub_dom_sub_breakup_effect = {
					DOMINANT = scope:actor
					SUBMISSIVE = scope:recipient
				}
				stress_impact = {
					lustful = minor_stress_impact_gain
					compassionate = minor_stress_impact_gain
				}
			}
		}
	}
}

# Regain submissive relation
sub_regain_submissive_interaction = {
	category = interaction_category_friendly
	common_interaction = yes

	auto_accept = yes
	ignores_pending_interaction_block = yes

	desc = sub_regain_submissive_interaction_desc

	is_shown = {
		scope:recipient = {
			can_set_relation_dominant_trigger = { CHARACTER = scope:actor }
			sub_has_obedience_modifier_trigger = { DOMINANT = scope:actor }
		}
	}

	on_accept = {
		scope:actor = {
			send_interface_message = {
				type = event_generic_neutral
				title = sub_regain_submissive_interaction_notification
				right_icon = scope:recipient
				set_relation_submissive = scope:recipient
				stress_impact = {
					lustful = minor_stress_impact_loss
				}
			}
		}
	}
}

#######
### TRAINING SCHEMES
#######

# Start sub_train_obedience scheme
sub_train_obedience_interaction = {
	category = interaction_category_friendly

	common_interaction = yes
	interface_priority = 150

	auto_accept = yes
	ignores_pending_interaction_block = yes

	scheme = sub_train_obedience

	send_name = START_SUB_TRAINING

	desc = sub_train_obedience_interaction_desc

	is_shown = {
		scope:recipient = {
			sub_can_be_obedience_trained_trigger = { DOMINANT = scope:actor }
		}
	}

	is_valid = {
		scope:actor = { is_ai = no }
	}

	is_valid_showing_failures_only = {
		scope:recipient = {
			sub_is_interested_in_obedience_training_as_submissive_of_trigger = { DOMINANT = scope:actor }
		}
		scope:actor = {
			sub_is_interested_in_obedience_training_as_dominant_of_trigger = { SUBMISSIVE = scope:recipient }
			can_start_scheme = {
				type = sway
				target = scope:recipient
			}
		}
	}

	on_accept = {
		scope:actor = {
			hidden_effect = {
				send_interface_toast = {
					title = sub_train_obedience_interaction_toast
					left_icon = scope:actor
					right_icon = scope:recipient

					start_scheme = {
						target = scope:recipient
						type = sub_train_obedience
					}
				}
			}
		}
	}
}

# Gain loyalty of already trained character
sub_train_obedience_accelerated_interaction = {
	category = interaction_category_friendly

	common_interaction = yes
	interface_priority = 150

	auto_accept = yes
	ignores_pending_interaction_block = yes

	desc = sub_train_obedience_accelerated_interaction_desc

	is_shown = {
		scope:recipient = {
			sub_can_be_accelerated_obedience_trained_trigger = { DOMINANT = scope:actor }
		}
	}

	is_valid = {
		scope:actor = { is_ai = no }
	}

	on_accept = {
		sub_train_obedience_success_effect = {
			DOMINANT = scope:actor
			SUBMISSIVE = scope:recipient
		}
		scope:actor = {
			hidden_effect = {
				send_interface_toast = {
					title = sub_train_obedience_accelerated_interaction_toast
					left_icon = scope:actor
					right_icon = scope:recipient
				}
			}
		}
	}
}

# Become regent of submissive liege
sub_dominant_demand_regency_interaction = {
	category = interaction_category_diarch

	auto_accept = yes
	ignores_pending_interaction_block = yes

	desc = sub_dominant_demand_regency_interaction_desc

	is_shown = {
		# The obvious.
		NOT = { scope:actor = scope:recipient }
		# Only useable on your liege.
		scope:recipient = scope:actor.liege

		scope:recipient = {
			# And said-liege can't be in a regency already...
			NOT = { has_active_diarchy = yes }
			# ... and must be subservient
			has_trait = sub_subservient_4
		}

		scope:actor = {
			# ... and must be your submissive
			has_relation_submissive = scope:recipient
		}

	}

	cooldown_against_recipient = { years = 1 }

	on_accept = {
		scope:recipient = {
			designate_diarch = scope:actor
			start_diarchy = regency
		}
	}
}

# Become liege of submissive
sub_dominant_demand_vassalization_interaction = {
	category = interaction_category_vassal
	icon = icon_vassal

	notification_text = SUB_DOMINANT_DEMAND_VASSALIZATION

	#interface_priority = 20
	#auto_accept = yes
	#ignores_pending_interaction_block = yes

	desc = sub_dominant_demand_vassalization_interaction_desc

	is_shown = {
		# The obvious.
		NOT = { scope:actor = scope:recipient }
		scope:recipient = {
			is_landed = yes
			#is_independent_ruler = yes
			#highest_held_title_tier < scope:actor.highest_held_title_tier
			#NOT = { government_has_flag = can_not_be_vassal_or_liege }
			# ... must be subservient
			has_trait = sub_subservient_4
		}

		scope:actor = {
			# ... and must be your submissive
			has_relation_submissive = scope:recipient
		}
	}

	cooldown_against_recipient = { years = 1 }

	#on_accept = {
	#	offer_vassalization_interaction_effect = yes
#		scope:actor = {
			#trigger_event = char_interaction.0001
		#}
	#}

	#is_valid_showing_failures_only = {
	#	scope:recipient = {
	#		is_at_war = no

	#		custom_description = {
	#			text = was_recently_granted_independence
	#			NOT = {
	#				has_opinion_modifier = {
	#					modifier = granted_independence_opinion
	#					target = scope:actor
	#				}
	#			}
	#		}
	#	}
	#}

	greeting = positive
	notification_text = SUB_DOMINANT_DEMAND_VASSALIZATION_INTERACTION_NOTIFICATION

	ai_min_reply_days = 5
	ai_max_reply_days = 10

	# High starting obligations
	send_option = {
		is_shown = {
			scope:recipient = {
				government_has_flag = government_is_feudal
			}
		}
		is_valid = {
			always = yes
		}
		flag = high_obligations								# If selected then scope:flag_name will be set to yes
		localization = high_obligations					# Loc_key for option label
	}

	# Nothing, for clans
	send_option = {
		is_shown = {
			scope:recipient = {
				government_has_flag = government_is_clan
			}
		}
		is_valid = {
			always = yes
		}
		flag = no_exemption_clan
		localization = no_exemption
	}

	send_options_exclusive = yes

	on_accept = {
		offer_vassalization_interaction_effect = yes
		scope:actor = {
			trigger_event = char_interaction.0001
		}
	}

	on_decline = {
		# On decline some response to your sub would be in place...
		scope:actor = {
			send_interface_toast = {
				title = sub_dominant_demand_vassalization_interaction.tt.failure
				left_icon = scope:recipient
				# AI lose opinion of their opposite number.
				if = {
					limit = { is_ai = yes }
					add_opinion = {
						target = scope:recipient
						modifier = insult_opinion
						opinion = -20
					}
				}
			}
		}
	}

	ai_accept = {
		# Try to make it 0 for most interactions
		base = 0

		modifier = { # I am independent and I do as I want
			scope:recipient = {
				is_independent_ruler = true
			}
			add = 50
		}

		# But what about my liege
		# modifier = {
		# 	add = {
		# 		if = {
		# 			limit = {
		# 				is_independent_ruler = false
		# 			}
		# 			opinion_modifier = {
		# 				who = scope:recipient
		# 				opinion_target = scope:recipient.liege
		# 				multiplier = -0.5
		# 			}
		# 		}
		# 	}
		# }

		# Do I even like my dominant?
		opinion_modifier = {
			who = scope:recipient
			opinion_target = scope:actor
			multiplier = 0.5
		}

		# Personality.
		ai_value_modifier = {
			# Meek characters are more likely to accept this.
			ai_boldness = -0.5
			# Characters with low energy just want to avoid all this planning.
			ai_energy = -0.25
		}

		# Relationships.
		## Positive.
		modifier = {
			add = 25
			desc = sub_dominant_demand_vassalization_interaction.tt.actor_is.positive_relationship
			has_any_good_relationship_with_character_trigger = { CHARACTER = scope:actor }
		}
		## Negative.
		modifier = {
			add = -50
			desc = sub_dominant_demand_vassalization_interaction.tt.actor_is.negative_relationship
			has_any_bad_relationship_with_character_trigger = { CHARACTER = scope:actor }
		}
	}
}

# TBD: Add loc and actual options, then uncomment is_shown and is_valid_showing_failures_only
# Spawn window to choose training
# sub_set_training_interaction = {
# 	category = interaction_category_friendly
# 	common_interaction = yes
# 	interface_priority = 95

# 	desc = "test desc"

# 	auto_accept = yes

# 	is_shown = {
# 		scope:recipient = {
# 			AND = {
# 				is_male = yes
# 				is_female = yes
# 			}
# 		}
# 	}
# 	# is_shown = {
# 	# 	scope:recipient = {
# 	# 		AND = {
# 	# 			is_adult = yes
# 	# 			is_female = yes
# 	# 			OR = {
# 	# 				is_consort_of = scope:actor
# 	# 				has_relation_lover = scope:actor
# 	# 				has_relation_soulmate = scope:actor
# 	# 				has_relation_dominant = scope:actor
# 	# 				has_opinion_modifier = {
# 	# 					target = scope:actor
# 	# 					modifier = sub_subservient_to_master
# 	# 				}
# 	# 			}
# 	# 		}
# 	# 	}
# 	# }

# 	# TBD: if no valid training options, show failure
# 	# is_valid_showing_failures_only = {
# 	# 	scope:recipient = { has_trait = sub_subservient_4 }
# 	# }

# 	on_accept = {
# 		scope:actor = {
# 			trigger_event = {
# 				on_action = sub_set_training
# 			}
# 		}
# 	}
# }

#######
### EXPERIMENTAL
#######

# Breast Enhancement
# sub_make_busty_interaction = {
# 	category = interaction_category_friendly
# 	common_interaction = yes
# 	interface_priority = 100

#   	desc = sub_make_busty_interaction_desc

#   	auto_accept = yes

# 	is_shown = {
# 		scope:recipient = {
# 			AND = {
# 				NOT = { has_trait = tits_big_3 }
# 				is_female = yes
# 				OR = {
# 					sub_has_sub_relation_trigger = { DOMINANT = scope:actor }
# 					sub_has_obedience_modifier_trigger = { DOMINANT = scope:actor }
# 				}
# 			}
# 		}
# 	}

# 	is_valid_showing_failures_only = {
# 		scope:recipient = { is_adult = yes }
# 	}

# 	on_accept = {
# 		scope:recipient = { add_trait = tits_big_3 }

# 		scope:actor = {
# 			send_interface_toast = {
# 				title = sub_breast_enhancement_order
# 				right_icon = scope:recipient
# 			}
# 		}
# 	}
# }


code:
SubMasterMistress = {
	type = character

	text = {
		trigger = {
			is_male = yes
		}
		localization_key = SUB_MASTER
	}

	text = {
		trigger = {
			is_female = yes
		}
		localization_key = SUB_MISTRESS
	}
}

SubMasterMistressPossessive = {
	type = character

	text = {
		trigger = {
			is_male = yes
		}
		localization_key = SUB_MASTER_POSSESSIVE
	}

	text = {
		trigger = {
			is_female = yes
		}
		localization_key = SUB_MISTRESS_POSSESSIVE
	}
}

SubGenderedGenitals = {
	type = character

	text = {
		trigger = {
			OR = {
				is_male = yes
				carn_is_futa_trigger = yes
			}
		}
		localization_key = SUB_PENIS
	}

	text = {
		trigger = {
			is_female = yes
			carn_is_futa_trigger = no
		}
		localization_key = SUB_PUSSY
	}
}

SubGenderedTorso = {
	type = character

	text = {
		trigger = {
			is_male = yes
		}
		localization_key = SUB_MALE_TORSO
	}

	text = {
		trigger = {
			is_female = yes
		}
		localization_key = SUB_FEMALE_TORSO
	}
}


code:
sub_train_obedience_scheme = {
	alias = { sub_train_obedience_schemes sub_train_obedience sub_train_obedience_action sub_train_obedience_scheme_single }
	parent = scheme
}

sub_subservience_concept = {
	alias = { sub_subservient_concept }
	texture = "gfx/interface/icons/traits/sub_subservient_4.dds"
}


code:
sub_valid_subs = {
	default = sub_valid_subs_all_genders
	sub_valid_subs_all_genders = {}
	sub_valid_subs_female_only = {}
	sub_valid_subs_male_only = {}
}

sub_gender_attraction_requirement = {
	default = sub_gender_attraction_requirement_enabled
	sub_gender_attraction_requirement_enabled = {}
	sub_gender_attraction_requirement_disabled = {}
}

sub_incestuous_training = {
	default = sub_incestuous_training_disabled
	sub_incestuous_training_enabled = {}
	sub_incestuous_training_disabled = {}
}


code:
action_can_obedience_train = {
	combine_into_one = yes

	check_create_action = {
		every_consort = {
			limit = {
				sub_can_be_obedience_trained_trigger = { DOMINANT = root }
				NOT = {
					any_targeting_scheme = {
						scheme_owner = root
					}
				}
			}
			try_create_important_action = {
				important_action_type = action_can_obedience_train
				actor = root
				recipient = this
			}
		}
		every_relation = {
			type = lover
			limit = {
				sub_can_be_obedience_trained_trigger = { DOMINANT = root }
				NOT = {
					any_targeting_scheme = {
						scheme_owner = root
					}
				}
			}
			try_create_important_action = {
				important_action_type = action_can_obedience_train
				actor = root
				recipient = this
			}
		}
		every_relation = {
			type = soulmate
			limit = {
				sub_can_be_obedience_trained_trigger = { DOMINANT = root }
				NOT = {
					any_targeting_scheme = {
						scheme_owner = root
					}
				}
			}
			try_create_important_action = {
				important_action_type = action_can_obedience_train
				actor = root
				recipient = this
			}
		}
	}

	effect = {
		open_interaction_window = {
			interaction = sub_train_obedience_interaction
			actor = scope:actor
			recipient = scope:recipient
		}
	}
}

action_can_accelerated_obedience_train = {
	combine_into_one = yes

	check_create_action = {
		every_consort = {
			limit = {
				sub_can_be_accelerated_obedience_trained_trigger = { DOMINANT = root }
			}
			try_create_important_action = {
				important_action_type = action_can_accelerated_obedience_train
				actor = root
				recipient = this
			}
		}
		every_relation = {
			type = lover
			limit = {
				sub_can_be_accelerated_obedience_trained_trigger = { DOMINANT = root }
			}
			try_create_important_action = {
				important_action_type = action_can_accelerated_obedience_train
				actor = root
				recipient = this
			}
		}
		every_relation = {
			type = soulmate
			limit = {
				sub_can_be_accelerated_obedience_trained_trigger = { DOMINANT = root }
			}
			try_create_important_action = {
				important_action_type = action_can_accelerated_obedience_train
				actor = root
				recipient = this
			}
		}
	}

	effect = {
		open_interaction_window = {
			interaction = sub_train_obedience_accelerated_interaction
			actor = scope:actor
			recipient = scope:recipient
		}
	}
}


code:
sub_recent_training_modifier = {
	icon = social_positive
	direct_vassal_opinion = 5
	monthly_prestige = 0.5
	attraction_opinion = 20
	stacking = yes
}

sub_training_milestone_scheme_power_1 = {
	icon = intrigue_positive
	scheme_power = 20
}

sub_training_milestone_scheme_power_2 = {
	icon = intrigue_positive
	scheme_power = 20
}

sub_training_milestone_scheme_power_3 = {
	icon = intrigue_positive
	scheme_power = 20
}

sub_positive_trait_scheme_power = {
	icon = intrigue_positive
	scheme_power = 20
	stacking = yes
}

sub_negative_trait_scheme_resistance = {
	icon = intrigue_negative
	scheme_resistance = 20
	stacking = yes
}

sub_highest_pos_opinion_scheme_power = {
	icon = intrigue_positive
	scheme_power = 40
}

sub_high_pos_opinion_scheme_power = {
	icon = intrigue_positive
	scheme_power = 30
}

sub_medium_pos_opinion_scheme_power = {
	icon = intrigue_positive
	scheme_power = 20
}

sub_low_pos_opinion_scheme_power = {
	icon = intrigue_positive
	scheme_power = 10
}

sub_highest_neg_opinion_scheme_resistance = {
	icon = intrigue_negative
	scheme_resistance = 40
}

sub_high_neg_opinion_scheme_resistance = {
	icon = intrigue_negative
	scheme_resistance = 30
}

sub_medium_neg_opinion_scheme_resistance = {
	icon = intrigue_negative
	scheme_resistance = 20
}

sub_low_neg_opinion_scheme_resistance = {
	icon = intrigue_negative
	scheme_resistance = 10
}


code:
sub_ask_to_pose = {
  events = {
    sub_pose.0001
  }
}

sub_set_clothing_rules = {
  events = {
    sub_rules.0002
  }
}

sub_set_chastity_rules = {
  events = {
    sub_rules.0001
  }
}

# sub_set_training = {
#   events = {
#     sub_training.0001
#   }
# }


code:
sub_train_obedience_start = {
	trigger = {
		exists = scope:scheme
	}

	first_valid = {
		sub_train_obedience_outcome.1000
	}
}

# Select target's response based on traits and opinion of player
sub_train_obedience_start_response = {
	first_valid = {
		sub_train_obedience_outcome.1001
		sub_train_obedience_outcome.1002
		sub_train_obedience_outcome.1003
	}
}

sub_train_obedience_milestone_1 = {
	events = {
		sub_train_obedience_outcome.1500
	}
}

sub_train_obedience_milestone_2= {
	events = {
		sub_train_obedience_outcome.1501
	}
}

sub_train_obedience_milestone_3 = {
	events = {
		sub_train_obedience_outcome.1502
	}
}


code:
sub_submissive_stress_reduce_scene = {
	trigger = {
		is_ai = no
		any_relation = {
			type = submissive
			exists = this
		}
	}

	effect = {
		ordered_relation = {
			type = submissive
			order_by = { value = 0 subtract = age } # prefer younger subs
			save_scope_as = lucky_submissive
		}
		trigger_event = {
			on_action = sub_stess_reduction_scene_selector
		}
	}
}

yearly_playable_pulse = {
	on_actions = {
		sub_submissive_stress_reduce_scene
	}
}

sub_stess_reduction_scene_selector = {
	random_events = {	# A single event will be picked to fire

		chance_to_happen = 40 # A percentage chance determining whether the events involved will be evaluated at all

		100 = sub_yearly.1001 	# The number is the weight for picking a specific event. The weight is factored by the event's weight_multiplier entry. (If no weight_multiplier is defined for the event, it is 1)
		# 200 = event_id_2
		# 100 = 0		# Having a "0" entry means that there is a chance no event fires, even if there are other valid events. Good for making sure that rare events don't always fire just because every other possible event is invalid.
	}
}


code:
sub_subservient_to_master = {
	opinion = 100
	years = 200
}

sub_made_me_go_naked = {
	opinion = -20
	years = 1
}

sub_positive_training_start = {
	opinion = 10
	months = 6
}

sub_negative_training_start = {
	opinion = -10
	months = 6
}


code:
sub_train_obedience = {

	skill = martial
	uses_agents = no
	power_per_skill_point = 2
	resistance_per_skill_point = 0
	minimum_progress_chance = 25
	maximum_progress_chance = 99
	maximum_success = 100

	# Does not use Spymaster
	spymaster_power_per_skill_point = 0
	spymaster_resistance_per_skill_point = 0

	allow = {
		is_ai = no
		scope:target = {
			sub_can_be_obedience_trained_trigger = { DOMINANT = scope:owner }
		}
	}

	valid = {
		is_imprisoned = no

		scope:target = {
			OR = {
				exists = location
				in_diplomatic_range = scope:owner
			}
		}
	}

	base_success_chance = {
		base = 1000
	}

	# Just fire conclusion event; no chance of discovery and scheme owner is always player
	on_ready = {
		scheme_owner = {
			trigger_event = sub_train_obedience_outcome.2000
		}
	}

	on_invalidated = {
		# Dead target
		if = {
			limit = {
				scheme_target = { is_alive = no }
			}
			scheme_owner = {
				trigger_event = sub_train_obedience_outcome.0002
			}
		}

		# Has none of the original qualifying relationships
		if = {
			limit = {
				NOT = { scheme_progress = 100 }
				scheme_target = { is_alive = yes }
				scheme_owner = {
					NOR = {
						is_consort_of = scope:scheme.scheme_target
						has_relation_lover = scope:scheme.scheme_target
						has_relation_soulmate = scope:scheme.scheme_target
						has_relation_submissive = scope:scheme.scheme_target
					}
				}
			}
			scheme_owner = {
				trigger_event = sub_train_obedience_outcome.0001
			}
		}

		# Misc conditions like being imprisoned, no more diplo range
		else_if = { #fallback invalidation
			limit = {
				OR = {
					scope:owner = { is_imprisoned = yes }

					scope:target = {
						NOR = {
							exists = location
							in_diplomatic_range = scope:owner
						}
					}
				}
			}
			scheme_owner = {
				send_interface_message = {
					type = event_generic_neutral
					title = sub_train_obedience
					left_icon = scope:target
					custom_tooltip = sub_train_obedience_invalid
				}
			}
		}
	}

	# trigger training start event
	on_start = {
		save_scope_as = scheme
		scheme_owner = { save_scope_as = owner }
		scheme_target = { save_scope_as = target }

		scope:scheme = {
			sub_train_obedience_set_dummy_variables = yes
			sub_set_scheme_opinion_modifier = yes
		}

		scope:owner = {
			trigger_event = {
				on_action = sub_train_obedience_start
				days = { 3 7 }
			}
		}

	}

	on_monthly = {
		save_scope_as = scheme
		scheme_owner = { save_scope_as = owner }
		scheme_target = { save_scope_as = target }

		scope:scheme = {
			sub_set_scheme_opinion_modifier = yes
		}

		# Milestone 1 event
		if = {
			limit = {
				scheme_progress >= 2
				NOT = { has_scheme_modifier = sub_training_milestone_scheme_power_1 }
			}
			scheme_owner = {
				trigger_event = { on_action = sub_train_obedience_milestone_1 }
			}
		}

		# Milestone 2 event
		if = {
			limit = {
				scheme_progress >= 4
				NOT = { has_scheme_modifier = sub_training_milestone_scheme_power_2 }
			}
			scheme_owner = {
				trigger_event = { on_action = sub_train_obedience_milestone_2 }
			}
		}
		# Milestone 3 event
		if = {
			limit = {
				scheme_progress >= 7
				NOT = { has_scheme_modifier = sub_training_milestone_scheme_power_3 }
			}
			scheme_owner = {
				trigger_event = { on_action = sub_train_obedience_milestone_3 }
			}
		}

	}
	success_desc = "SUB_TRAIN_OBEDIENCE_SUCCESS_DESC"
}


code:
sub_train_obedience_success_effect = {
	$SUBMISSIVE$ = {
		hidden_effect = {
			remove_trait = sub_subservient_3
		}
		add_trait = sub_subservient_4
		add_opinion = {
			modifier = sub_subservient_to_master
			target = $DOMINANT$
		}
	}
	$DOMINANT$ = {
		add_hook = {
			type = loyalty_hook
			target = $SUBMISSIVE$
		}
		if = {
			limit = { can_set_relation_submissive_trigger = { CHARACTER = $SUBMISSIVE$ } }
			set_relation_submissive = $SUBMISSIVE$
		}
		add_character_modifier = {
			modifier = sub_recent_training_modifier
			years = 5
		}
	}
}


code:
sub_dom_sub_breakup_effect = {
	$DOMINANT$ = {
		remove_relation_submissive = $SUBMISSIVE$
	}
}