sz-orm-vector 10.13.4

//! SZ-ORM pgvector Extension
//!
//! Provides PostgreSQL pgvector vector similarity search capabilities, supporting three implementations:
//!
//! - **In-memory implementation** (`InMemoryVectorStore`): Pure Rust vector computation, no database connection, suitable for testing and benchmarking
//! - **Stub implementation** (`StubVectorStore`): All methods return Unsupported, suitable for debug placeholder
//! - **Real implementation** (`RealPgVectorStore`, requires `real-pg` feature): Connects to PostgreSQL + pgvector via tokio-postgres
//!
//! # Quick Start
//!
//! ```rust
//! use sz_orm_vector::{InMemoryVectorStore, PgVectorStore, VectorRecord, VectorMetric};
//!
//! # #[tokio::main]
//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
//! let store = InMemoryVectorStore::new();
//! store.create_collection("docs", 3, None).await?;
//!
//! let record = VectorRecord::new("doc1", vec![1.0, 0.0, 0.0]);
//! store.insert("docs", vec![record]).await?;
//!
//! let results = store.search("docs", &[1.0, 0.0, 0.0], 5).await?;
//! println!("found {} results", results.len());
//! # Ok(())
//! # }
//! ```

pub mod error;
pub mod extensions;
pub mod memory;
pub mod stub;

#[cfg(feature = "embedding-store")]
pub mod embedding_store;

#[cfg(feature = "real-pg")]
pub mod real_pg;

// v4.0.0 M3：混合搜索（hybrid-search feature gate 隔离）
#[cfg(feature = "hybrid-search")]
pub mod hybrid_search;

pub use error::VectorError;
pub use extensions::{
    AnnIndexDef, AnnIndexRegistry, AnnIndexType, BatchOpsExt, DimensionValidator, HnswParams,
    IvfflatParams, MemoryBatchOps, SimilarityAlgorithms, VectorNormalizer, MAX_VECTOR_DIMENSION,
    MIN_VECTOR_DIMENSION,
};
pub use memory::InMemoryVectorStore;
pub use stub::StubVectorStore;

#[cfg(feature = "real-pg")]
pub use real_pg::{RealPgConfig, RealPgVectorStore};

use async_trait::async_trait;
use std::collections::HashMap;
use std::str::FromStr;

/// M-16 fix: top_k maximum limit
///
/// Limit top_k upper bound to prevent:
/// - Large k values causing memory explosion (each SearchResult contains full vector)
/// - Performance issues with database/vector engine executing huge k queries
/// - Malicious caller triggering OOM via top_k=usize::MAX
pub const MAX_TOP_K: usize = 10_000;

/// M-16 fix: Validate whether top_k is within reasonable range
///
/// - `top_k = 0`: Returns `TopKExceeded` error (meaningless query)
/// - `top_k > MAX_TOP_K`: Returns `TopKExceeded` error
/// - `1 <= top_k <= MAX_TOP_K`: Returns Ok
pub fn validate_top_k(top_k: usize) -> Result<usize, VectorError> {
    if top_k == 0 {
        return Err(VectorError::TopKExceeded {
            requested: top_k,
            max: MAX_TOP_K,
        });
    }
    if top_k > MAX_TOP_K {
        return Err(VectorError::TopKExceeded {
            requested: top_k,
            max: MAX_TOP_K,
        });
    }
    Ok(top_k)
}

/// Vector record
#[derive(Debug, Clone)]
pub struct VectorRecord {
    pub id: String,
    pub vector: Vec<f32>,
    pub score: Option<f32>,
    pub metadata: Option<HashMap<String, serde_json::Value>>,
}

impl VectorRecord {
    pub fn new(id: impl Into<String>, vector: Vec<f32>) -> Self {
        Self {
            id: id.into(),
            vector,
            score: None,
            metadata: None,
        }
    }

    pub fn with_score(mut self, score: f32) -> Self {
        self.score = Some(score);
        self
    }

    pub fn with_metadata(mut self, metadata: HashMap<String, serde_json::Value>) -> Self {
        self.metadata = Some(metadata);
        self
    }
}

/// Search result
#[derive(Debug, Clone)]
pub struct SearchResult {
    pub id: String,
    pub score: f32,
    pub vector: Vec<f32>,
    pub text: Option<String>,
    pub metadata: Option<HashMap<String, serde_json::Value>>,
}

impl SearchResult {
    pub fn new(id: impl Into<String>, score: f32, vector: Vec<f32>) -> Self {
        Self {
            id: id.into(),
            score,
            vector,
            text: None,
            metadata: None,
        }
    }

    pub fn with_text(mut self, text: impl Into<String>) -> Self {
        self.text = Some(text.into());
        self
    }

    pub fn with_metadata(mut self, metadata: HashMap<String, serde_json::Value>) -> Self {
        self.metadata = Some(metadata);
        self
    }
}

/// Vector distance metric
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub enum VectorMetric {
    #[default]
    Cosine,
    Euclidean,
    DotProduct,
}

impl VectorMetric {
    /// pgvector operator mapping
    pub fn pg_operator(&self) -> &'static str {
        match self {
            VectorMetric::Cosine => "<=>",
            VectorMetric::Euclidean => "<->",
            VectorMetric::DotProduct => "<#>",
        }
    }

    pub fn as_str(&self) -> &'static str {
        match self {
            VectorMetric::Cosine => "cosine",
            VectorMetric::Euclidean => "euclidean",
            VectorMetric::DotProduct => "dotproduct",
        }
    }
}

impl FromStr for VectorMetric {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "cosine" => Ok(VectorMetric::Cosine),
            "euclidean" => Ok(VectorMetric::Euclidean),
            "dotproduct" => Ok(VectorMetric::DotProduct),
            _ => Err(format!("unknown vector metric: {}", s)),
        }
    }
}

/// Vector Store core trait
///
/// Provides CRUD and similarity search capabilities for vector collections.
/// All methods are async, suitable for real database I/O.
#[async_trait]
pub trait PgVectorStore: Send + Sync {
    /// Create collection
    async fn create_collection(
        &self,
        name: &str,
        dimension: usize,
        metric: Option<VectorMetric>,
    ) -> Result<(), VectorError>;

    /// Delete collection
    async fn delete_collection(&self, name: &str) -> Result<(), VectorError>;

    /// Insert vector record (upsert semantics: same id overwrites)
    async fn insert(&self, collection: &str, records: Vec<VectorRecord>)
        -> Result<(), VectorError>;

    /// Similarity search
    ///
    /// M-16 fix: `top_k` must be in `[1, MAX_TOP_K]` range.
    /// Implementations should call `validate_top_k(top_k)?` before executing search.
    async fn search(
        &self,
        collection: &str,
        query: &[f32],
        top_k: usize,
    ) -> Result<Vec<SearchResult>, VectorError>;

    /// Get single record
    async fn get(&self, collection: &str, id: &str) -> Result<Option<VectorRecord>, VectorError>;

    /// Delete record
    async fn delete(&self, collection: &str, ids: Vec<String>) -> Result<u64, VectorError>;

    /// Count records
    async fn count(&self, collection: &str) -> Result<usize, VectorError>;
}
// v7.3.0 任务 3.5：向量 ANN 加速与召回指标与租户隔离
// 启用 ann-accel feature 时编译，默认不启用

/// ANN 搜索结果（含召回率指标）
#[derive(Debug, Clone)]
pub struct AnnSearchResult {
    /// 检索结果
    pub records: Vec<SearchResult>,
    /// 召回率（ANN 结果 vs 暴力 kNN，0.0 ~ 1.0）
    pub recall_rate: f64,
    /// 延迟（毫秒）
    pub latency_ms: u64,
}

/// ANN 加速检索 trait（继承 PgVectorStore）
///
/// 提供 ANN 索引加速检索能力，支持租户隔离过滤和召回率指标。
/// 内存实现使用精确 kNN（召回率 = 1.0），真实 ANN 加速需启用 pgvector HNSW/IVF 索引。
#[cfg(feature = "ann-accel")]
#[async_trait]
pub trait AnnAccelerated: PgVectorStore {
    /// ANN 加速检索
    ///
    /// # 参数
    /// - `collection`: 集合名称
    /// - `query`: 查询向量
    /// - `top_k`: 返回结果数
    /// - `tenant_id`: 租户 ID（可选，用于租户隔离过滤）
    ///
    /// # 返回值
    /// - `Ok(AnnSearchResult)`: 检索结果 + 召回率 + 延迟
    /// - `Err(VectorError)`: 检索失败
    async fn ann_search(
        &self,
        collection: &str,
        query: &[f32],
        top_k: usize,
        tenant_id: Option<&str>,
    ) -> Result<AnnSearchResult, VectorError>;
}

#[cfg(feature = "ann-accel")]
#[async_trait]
impl AnnAccelerated for InMemoryVectorStore {
    async fn ann_search(
        &self,
        collection: &str,
        query: &[f32],
        top_k: usize,
        tenant_id: Option<&str>,
    ) -> Result<AnnSearchResult, VectorError> {
        use std::time::Instant;
        let start = Instant::now();

        // 精确 kNN 检索（内存实现中 ANN = kNN，召回率 = 1.0）
        let mut results = self.search(collection, query, top_k).await?;

        // 租户隔离过滤：只保留 metadata 中 tenant_id 匹配的记录
        if let Some(tid) = tenant_id {
            results.retain(|r| {
                r.metadata
                    .as_ref()
                    .and_then(|m| m.get("tenant_id"))
                    .and_then(|v| v.as_str())
                    == Some(tid)
            });
        }

        let latency_ms = start.elapsed().as_millis() as u64;

        // 召回率 = 1.0（内存实现使用精确 kNN）
        let recall_rate = 1.0;

        Ok(AnnSearchResult {
            records: results,
            recall_rate,
            latency_ms,
        })
    }
}