data-sampler

/**
 * @file types.ts
 * @description Plugin-specific type definitions for the data-sampler plugin.
 *
 * Conventions:
 * - All types are readonly where possible.
 * - SamplingStrategy enum covers the 3 supported strategies.
 * - Chunk represents a single contiguous segment of the source text.
 * - SampleResult is the top-level output of any sampling operation.
 */

// ---------------------------------------------------------------------------
// Strategy enumeration
// ---------------------------------------------------------------------------

/** The 3 supported text sampling strategies. */
export type SamplingStrategy = "priority" | "relevance" | "length";

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

/**
 * Configuration for a single sampling operation.
 *
 * @property strategy - Which sampling strategy to use.
 * @property chunkSize - Target character size per chunk (default: 2000).
 * @property overlap - Number of overlapping characters between chunks (default: 0).
 * @property maxChunks - Maximum number of chunks to return (default: unlimited).
 * @property query - Required for "relevance" strategy; used for BM25-style scoring.
 */
export interface SamplingConfig {
  readonly strategy: SamplingStrategy;
  readonly chunkSize?: number;
  readonly overlap?: number;
  readonly maxChunks?: number;
  readonly query?: string;
}

// ---------------------------------------------------------------------------
// Data structures
// ---------------------------------------------------------------------------

/** A single chunk of sampled text with metadata. */
export interface Chunk {
  /** 0-based index of this chunk in the result set. */
  readonly index: number;
  /** The chunk text content. */
  readonly text: string;
  /** Relevance / priority score (0-1, higher is better). */
  readonly score: number;
  /** Character offset in the original text where this chunk starts. */
  readonly startOffset: number;
  /** Character offset in the original text where this chunk ends (exclusive). */
  readonly endOffset: number;
}

/** The complete result of a sampling operation. */
export interface SampleResult {
  /** The sampled chunks, sorted per strategy semantics. */
  readonly chunks: readonly Chunk[];
  /** Total number of chunks produced (before maxChunks truncation). */
  readonly totalChunks: number;
  /** Length in characters of the original input text. */
  readonly originalLength: number;
  /** The strategy that was used for sampling. */
  readonly strategy: SamplingStrategy;
}

// ---------------------------------------------------------------------------
// Dataset loading (hf datasets server)
// ---------------------------------------------------------------------------

/**
 * Parameters for the load_dataset tool.
 *
 * @property dataset - HF dataset name (e.g. "squad", "imdb", "ibm/duorc").
 * @property split - Dataset split: "train", "test", or "validation" (default: "train").
 * @property config - Dataset config/subset name (required for multi-config datasets).
 * @property max_samples - Number of rows to fetch (default: 5, max: 100).
 * @property query - Optional BM25 search query text to filter rows.
 * @property format - Output format: "text", "json", or "prompt" (default: "text").
 * @property template - Optional template string with {text} placeholder for each sample.
 */
export interface LoadDatasetParams {
  readonly dataset: string;
  readonly split?: string;
  readonly config?: string;
  readonly max_samples?: number;
  readonly query?: string;
  readonly format?: "text" | "json" | "prompt" | "fewshot";
  readonly template?: string;
  /**
   * Optional comma/space-separated string of row indices to fetch.
   * E.g. "0,5,10" or "0 5 10". When provided, only those rows are returned.
   * The tool fetches a contiguous batch from min index to max index,
   * then filters to just the requested indices.
   */
  readonly indices?: string;
  /**
   * Columns to treat as input (for fewshot format). If not specified,
   * all columns except output_column are treated as input.
   */
  readonly input_columns?: readonly string[];
  /**
   * Column to treat as output/answer (for fewshot format). Required for fewshot mode.
   */
  readonly output_column?: string;
}

/** A single sample row from a HuggingFace dataset. */
export interface DatasetSample {
  /** 0-based index of this row in the dataset. */
  readonly index: number;
  /** The raw column/value data for this row. */
  readonly data: Readonly<Record<string, unknown>>;
}

/** The complete result of a load_dataset operation. */
export interface LoadDatasetResult {
  /** The HF dataset name that was queried. */
  readonly dataset: string;
  /** The split that was queried. */
  readonly split: string;
  /** The config/subset used, or null if auto-detected. */
  readonly config: string | null;
  /** The sampled rows from the dataset. */
  readonly samples: readonly DatasetSample[];
  /** Total number of rows available in the dataset split (not just returned). */
  readonly totalRows: number;
  /** Number of samples returned in this response. */
  readonly count: number;
  /** The format used for the formatted text output. */
  readonly format: string;
  /** Formatted text representation of all samples, per the requested format. */
  readonly formatted: string;
}