forge-embed

use forge_embed;

pub mod chunking;

pub mod document;

pub mod embed;

pub mod error;

pub mod rerank;

pub mod similarity;

pub mod vector_store;

pub mod prelude;

pub use crate::chunking::{RecursiveCharacterSplitter, TextSplitter, TokenSplitter};

pub use crate::document::{Document, DocumentLoader, JsonLoader, TextLoader};

pub use crate::embed::{EmbeddingProvider, EmbeddingResult};

pub use crate::error::{EmbedResult, ForgeEmbedError};

pub use crate::rerank::{RerankResult, Reranker};

pub use crate::similarity::{cosine_similarity, dot_product, euclidean_distance};

pub use crate::vector_store::{InMemoryVectorStore, SearchResult, VectorEntry, VectorStore};

pub trait TextSplitter: Send + Sync {
    /// Splits the input text into chunks.
    ///
    /// # Arguments
    ///
    /// * `text` - The text to split.
    ///
    /// # Returns
    ///
    /// A vector of non-empty string chunks. Returns an empty vector if the
    /// input is empty.
    fn split(&self, text: &str) -> Vec<String>;
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecursiveCharacterSplitter {

}

pub fn new(chunk_size: usize, overlap: usize) -> Self;

pub fn with_separators(chunk_size: usize, overlap: usize, separators: Vec<String>) -> Self;

pub fn chunk_size(&self) -> usize;

pub fn overlap(&self) -> usize;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenSplitter {

}

pub fn new(tokens_per_chunk: usize, overlap_tokens: usize) -> Self;

pub fn tokens_per_chunk(&self) -> usize;

pub fn overlap_tokens(&self) -> usize;

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Document {
/// The text content of the document.

pub content: String,
/// Arbitrary metadata associated with this document.

pub metadata: serde_json::Value
}

pub fn new(content: impl Into<String>) -> Self;

pub fn with_metadata(content: impl Into<String>, metadata: serde_json::Value) -> Self;

pub fn is_empty(&self) -> bool;

pub fn len(&self) -> usize;

pub trait DocumentLoader: Send + Sync {
    /// Loads documents from the given source string.
    ///
    /// # Arguments
    ///
    /// * `source` - The source to load from. Interpretation depends on the
    ///   implementation (raw text, JSON string, file path, etc.).
    ///
    /// # Returns
    ///
    /// A vector of loaded documents.
    ///
    /// # Errors
    ///
    /// Returns [`ForgeEmbedError`] if loading fails (e.g., invalid format,
    /// empty input).
    fn load(&self, source: &str) -> EmbedResult<Vec<Document>>;
}

#[derive(Debug, Clone, Copy, Default)]
pub struct TextLoader;

#[derive(Debug, Clone, Default)]
pub struct JsonLoader {

}

pub fn new(content_field: Option<String>) -> Self;

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct EmbeddingResult {
/// The embedding vector.

pub vector: Vec<f64>,
/// The model that produced this embedding.

pub model: String,
/// The number of dimensions in the embedding vector.

pub dimensions: usize
}

#[async_trait]
pub trait EmbeddingProvider: Send + Sync {
    /// Returns the model identifier for this provider (e.g., "text-embedding-3-small").
    fn model_id(&self) -> &str;

    /// Generates embedding vectors for one or more text inputs.
    ///
    /// # Arguments
    ///
    /// * `input` - A slice of strings to embed. Must contain at least one item.
    ///
    /// # Returns
    ///
    /// A vector of embedding vectors, one per input string, in the same order.
    ///
    /// # Errors
    ///
    /// Returns [`ForgeEmbedError::ModelError`] if the provider call fails,
    /// or [`ForgeEmbedError::EmptyInput`] if the input slice is empty.
    async fn embed(&self, input: &[String]) -> Result<Vec<Vec<f64>>, ForgeEmbedError>;
}

pub async fn embed(provider: &dyn EmbeddingProvider, text: &str) -> EmbedResult<EmbeddingResult>;

pub async fn embed_many(
    provider: &dyn EmbeddingProvider,
    texts: &[String],
) -> EmbedResult<Vec<EmbeddingResult>>;

#[derive(Debug, Error)]
pub enum ForgeEmbedError {
    /// The embedding model returned an error.
    #[error("embedding model '{model}' failed: {reason}")]
    ModelError {
        /// The model identifier that was used.
        model: String,
        /// A description of what went wrong.
        reason: String,
    },

    /// Vector dimensions do not match for the requested operation.
    #[error("dimension mismatch: vector A has {a} dimensions but vector B has {b} dimensions; both must be equal for {operation}")]
    DimensionMismatch {
        /// Dimensions of the first vector.
        a: usize,
        /// Dimensions of the second vector.
        b: usize,
        /// The operation that required matching dimensions.
        operation: String,
    },

    /// An empty input was provided where at least one item is required.
    #[error("empty input for '{operation}': at least one item is required")]
    EmptyInput {
        /// The operation that received empty input.
        operation: String,
    },

    /// A vector store operation failed.
    #[error("vector store error during '{operation}': {reason}")]
    StoreError {
        /// The store operation that failed (e.g., "insert", "search", "delete").
        operation: String,
        /// A description of what went wrong.
        reason: String,
    },

    /// A text chunking operation failed.
    #[error("chunking error: {reason}")]
    ChunkingError {
        /// A description of what went wrong during chunking.
        reason: String,
    },

    /// An error propagated from the `forge-core` crate.
    #[error(transparent)]
    Core(#[from] ForgeError),
}

pub type EmbedResult<T> = Result<T, ForgeEmbedError>;

pub mod chunking;

pub mod document;

pub mod embed;

pub mod error;

pub mod rerank;

pub mod similarity;

pub mod vector_store;

pub mod prelude;

pub use crate::chunking::{RecursiveCharacterSplitter, TextSplitter, TokenSplitter};

pub use crate::document::{Document, DocumentLoader, JsonLoader, TextLoader};

pub use crate::embed::{EmbeddingProvider, EmbeddingResult};

pub use crate::error::{EmbedResult, ForgeEmbedError};

pub use crate::rerank::{RerankResult, Reranker};

pub use crate::similarity::{cosine_similarity, dot_product, euclidean_distance};

pub use crate::vector_store::{InMemoryVectorStore, SearchResult, VectorEntry, VectorStore};

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RerankResult {
/// The original index of this document in the input slice.

pub index: usize,
/// The relevance score (higher is more relevant).

pub score: f64,
/// The document text.

pub document: String
}

#[async_trait]
pub trait Reranker: Send + Sync {
    /// Reranks documents by relevance to a query.
    ///
    /// # Arguments
    ///
    /// * `query` - The query string to rank documents against.
    /// * `documents` - The documents to rerank.
    /// * `top_k` - If provided, return only the top K most relevant results.
    ///
    /// # Returns
    ///
    /// A vector of [`RerankResult`] sorted by descending relevance score.
    ///
    /// # Errors
    ///
    /// Returns [`ForgeEmbedError`] if the reranking operation fails.
    async fn rerank(
        &self,
        query: &str,
        documents: &[String],
        top_k: Option<usize>,
    ) -> Result<Vec<RerankResult>, ForgeEmbedError>;
}

pub async fn rerank(
    provider: &dyn EmbeddingProvider,
    query: &str,
    documents: &[String],
    top_k: Option<usize>,
) -> EmbedResult<Vec<RerankResult>>;

pub fn cosine_similarity(a: &[f64], b: &[f64]) -> EmbedResult<f64>;

pub fn euclidean_distance(a: &[f64], b: &[f64]) -> EmbedResult<f64>;

pub fn dot_product(a: &[f64], b: &[f64]) -> EmbedResult<f64>;

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct VectorEntry {
/// Unique identifier for this entry.

pub id: String,
/// The embedding vector.

pub vector: Vec<f64>,
/// Arbitrary metadata associated with this entry.

pub metadata: serde_json::Value
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SearchResult {
/// The matched vector entry.

pub entry: VectorEntry,
/// The similarity score (higher is more similar).

pub score: f64
}

#[async_trait]
pub trait VectorStore: Send + Sync {
    /// Inserts or updates a vector entry in the store.
    ///
    /// If an entry with the same `id` already exists, it is replaced.
    ///
    /// # Arguments
    ///
    /// * `entry` - The vector entry to insert.
    ///
    /// # Errors
    ///
    /// Returns [`ForgeEmbedError::StoreError`] if the insertion fails.
    async fn insert(&mut self, entry: VectorEntry) -> EmbedResult<()>;

    /// Searches for the nearest vectors to the query.
    ///
    /// Returns up to `top_k` results sorted by descending similarity score.
    ///
    /// # Arguments
    ///
    /// * `query` - The query embedding vector.
    /// * `top_k` - Maximum number of results to return.
    ///
    /// # Returns
    ///
    /// A vector of [`SearchResult`] sorted by descending similarity score.
    ///
    /// # Errors
    ///
    /// Returns [`ForgeEmbedError::StoreError`] if the search fails.
    /// Returns [`ForgeEmbedError::EmptyInput`] if the query vector is empty.
    async fn search(&self, query: &[f64], top_k: usize) -> EmbedResult<Vec<SearchResult>>;

    /// Deletes a vector entry by ID.
    ///
    /// # Arguments
    ///
    /// * `id` - The identifier of the entry to delete.
    ///
    /// # Returns
    ///
    /// `true` if the entry was found and deleted, `false` if it did not exist.
    ///
    /// # Errors
    ///
    /// Returns [`ForgeEmbedError::StoreError`] if the deletion fails.
    async fn delete(&mut self, id: &str) -> EmbedResult<bool>;
}

#[derive(Debug, Default)]
pub struct InMemoryVectorStore {

}

pub fn new() -> Self;

pub fn len(&self) -> usize;

pub fn is_empty(&self) -> bool;

pub fn contains(&self, id: &str) -> bool;

pub fn get(&self, id: &str) -> Option<&VectorEntry>;