@forge-sdk/embed

import * as api from '@forge-sdk/embed';

export interface TextSplitter {
  /**
   * Splits text into chunks.
   *
   * @param text - The text to split.
   * @returns An array of text chunks.
   */
  split(text: string): string[];
}

export interface RecursiveCharacterSplitterOptions {
  /** Maximum number of characters per chunk. Defaults to 1000. */
  readonly chunkSize?: number;
  /** Number of overlapping characters between consecutive chunks. Defaults to 200. */
  readonly chunkOverlap?: number;
  /**
   * Separator hierarchy from coarsest to finest. Defaults to
   * `['\n\n', '\n', '. ', ' ', '']` (paragraphs, lines, sentences, words, chars).
   */
  readonly separators?: readonly string[];
}

export class RecursiveCharacterSplitter implements TextSplitter {
  constructor(options: RecursiveCharacterSplitterOptions = {});
  split(text: string): string[];
}

export interface TokenSplitterOptions {
  /** Maximum number of tokens per chunk. Defaults to 256. */
  readonly tokensPerChunk?: number;
  /** Number of overlapping tokens between consecutive chunks. Defaults to 0. */
  readonly tokenOverlap?: number;
}

export class TokenSplitter implements TextSplitter {
  constructor(options: TokenSplitterOptions = {});
  split(text: string): string[];
}

export interface Document {
  /** The text content of the document. */
  readonly content: string;
  /** Optional key-value metadata for filtering and context. */
  readonly metadata: Record<string, unknown>;
}

export interface DocumentLoader {
  /**
   * Loads documents from the given source.
   *
   * @param source - The source content (text, JSON string, URL, etc.).
   * @returns An array of loaded documents.
   * @throws {ForgeEmbedError} If loading fails.
   */
  load(source: string): Promise<Document[]>;
}

export class TextLoader implements DocumentLoader {
  async load(source: string): Promise<Document[]>;
}

export class JsonLoader implements DocumentLoader {
  constructor(contentField?: string);
  async load(source: string): Promise<Document[]>;
}

export interface EmbeddingResult {
  /** The embedding vector as an array of floating-point numbers. */
  readonly vector: readonly number[];
  /** The model identifier that produced this embedding. */
  readonly model: string;
  /** The dimensionality of the embedding vector. */
  readonly dimensions: number;
}

export interface EmbeddingProvider {
  /** Returns the model identifier (e.g., 'text-embedding-3-small'). */
  modelId(): string;

  /**
   * Embeds one or more text strings into vectors.
   *
   * @param texts - The text strings to embed.
   * @returns An array of EmbeddingResult, one per input text.
   * @throws {ForgeEmbedError} If the model fails or input is invalid.
   */
  embed(texts: string[]): Promise<EmbeddingResult[]>;
}

export async function embed(
  provider: EmbeddingProvider,
  text: string
): Promise<EmbeddingResult>;

export async function embedMany(
  provider: EmbeddingProvider,
  texts: string[]
): Promise<EmbeddingResult[]>;

export const ForgeEmbedErrorCode /* type inferred in source */;

export type ForgeEmbedErrorCodeType = (typeof ForgeEmbedErrorCode)[keyof typeof ForgeEmbedErrorCode];

export class ForgeEmbedError extends Error {
  public readonly code: ForgeEmbedErrorCodeType;
  static modelError(model: string, reason: string): ForgeEmbedError;
  static dimensionMismatch(expected: number, actual: number): ForgeEmbedError;
  static emptyInput(operation: string): ForgeEmbedError;
  static storeError(operation: string, reason: string): ForgeEmbedError;
  static chunkingError(reason: string): ForgeEmbedError;
  static core(reason: string): ForgeEmbedError;
}

export { ForgeEmbedError, ForgeEmbedErrorCode, type ForgeEmbedErrorCodeType } from './error.js';

export { type EmbeddingProvider, type EmbeddingResult, embed, embedMany } from './embed.js';

export { cosineSimilarity, euclideanDistance, dotProduct } from './similarity.js';

export { type Reranker, type RerankResult, rerank } from './rerank.js';

export {
  type VectorStore,
  type VectorEntry,
  type SearchResult,
  InMemoryVectorStore,
} from './vector-store.js';

export {
  type TextSplitter,
  RecursiveCharacterSplitter,
  type RecursiveCharacterSplitterOptions,
  TokenSplitter,
  type TokenSplitterOptions,
} from './chunking.js';

export {
  type Document,
  type DocumentLoader,
  TextLoader,
  JsonLoader,
} from './document.js';

export interface RerankResult {
  /** The original index of this document in the input array. */
  readonly index: number;
  /** The relevance score (higher is more relevant). */
  readonly score: number;
  /** The document text. */
  readonly document: string;
}

export interface Reranker {
  /**
   * Reranks documents by relevance to a query.
   *
   * @param query - The search query.
   * @param documents - The documents to rerank.
   * @param topK - Maximum number of results to return.
   * @returns Reranked documents sorted by score descending.
   */
  rerank(query: string, documents: string[], topK: number): Promise<RerankResult[]>;
}

export async function rerank(
  provider: EmbeddingProvider,
  query: string,
  documents: string[],
  topK: number
): Promise<RerankResult[]>;

export function cosineSimilarity(a: readonly number[], b: readonly number[]): number;

export function euclideanDistance(a: readonly number[], b: readonly number[]): number;

export function dotProduct(a: readonly number[], b: readonly number[]): number;

export interface VectorEntry {
  /** The unique identifier for this entry. */
  readonly id: string;
  /** The embedding vector. */
  readonly vector: readonly number[];
  /** Optional key-value metadata for filtering and retrieval. */
  readonly metadata: Record<string, unknown>;
}

export interface SearchResult {
  /** The matching vector entry. */
  readonly entry: VectorEntry;
  /** The similarity score (higher is more similar). */
  readonly score: number;
}

export interface VectorStore {
  /**
   * Inserts a vector entry into the store.
   *
   * If an entry with the same ID already exists, it is overwritten.
   *
   * @param entry - The vector entry to insert.
   * @throws {ForgeEmbedError} If the insert fails.
   */
  insert(entry: VectorEntry): Promise<void>;

  /**
   * Searches for the top K most similar vectors to the query.
   *
   * @param queryVector - The query vector.
   * @param topK - Maximum number of results to return.
   * @returns An array of SearchResult sorted by score descending.
   * @throws {ForgeEmbedError} If the search fails.
   */
  search(queryVector: readonly number[], topK: number): Promise<SearchResult[]>;

  /**
   * Deletes a vector entry by its identifier.
   *
   * @param id - The identifier of the entry to delete.
   * @returns `true` if the entry was found and deleted, `false` otherwise.
   * @throws {ForgeEmbedError} If the delete fails.
   */
  delete(id: string): Promise<boolean>;
}

export class InMemoryVectorStore implements VectorStore {
  async insert(entry: VectorEntry): Promise<void>;
  async search(queryVector: readonly number[], topK: number): Promise<SearchResult[]>;
  async delete(id: string): Promise<boolean>;
  get size(): number;
}