Skip to main content

ailake_core/
schema.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2use crate::types::{VectorMetric, VectorPrecision};
3use serde::{Deserialize, Serialize};
4
5/// Canonical column names for LLM-context tables.
6/// ContextAssembler reads columns by these names.
7pub mod llm_columns {
8    pub const CHUNK_ID: &str = "chunk_id";
9    pub const DOCUMENT_ID: &str = "document_id";
10    pub const CHUNK_INDEX: &str = "chunk_index";
11    pub const TOTAL_CHUNKS: &str = "total_chunks";
12    pub const CHUNK_TEXT: &str = "chunk_text";
13    pub const DOCUMENT_TITLE: &str = "document_title";
14    pub const SECTION_PATH: &str = "section_path";
15    pub const PRECEDING_CONTEXT: &str = "preceding_context";
16    pub const FOLLOWING_CONTEXT: &str = "following_context";
17    pub const DOCUMENT_SUMMARY: &str = "document_summary";
18    pub const CHUNK_SUMMARY: &str = "chunk_summary";
19    pub const SOURCE_URI: &str = "source_uri";
20    pub const PAGE_NUMBER: &str = "page_number";
21    pub const CREATED_AT: &str = "created_at";
22    pub const DOCUMENT_DATE: &str = "document_date";
23    pub const EMBEDDING: &str = "embedding";
24    pub const CONTEXT_EMBEDDING: &str = "context_embedding";
25}
26
27/// Vector storage configuration applied at table creation time.
28/// Stored in Iceberg metadata.json properties.
29#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct VectorStoragePolicy {
31    pub column_name: String,
32    pub dim: u32,
33    pub metric: VectorMetric,
34    pub precision: VectorPrecision,
35    pub pq: Option<PQConfig>,
36    pub keep_raw_for_reranking: bool,
37    /// Normalize each input vector to unit L2 length before indexing.
38    /// Enables the NormalizedCosine fast path in HNSW: distance = 1 - dot(a, b),
39    /// no sqrt, ~2× faster distance computation. Semantics unchanged — same top-k
40    /// results as Cosine. Most embedding models (OpenAI, Cohere, etc.) produce
41    /// nearly-unit vectors; enabling this adds negligible write overhead.
42    #[serde(default)]
43    pub pre_normalize: bool,
44    /// HNSW M parameter — connections per node. `None` = default (16).
45    /// Higher M → better recall, more memory, slower build.
46    /// Recommended values: 8 (low-memory), 16 (default), 32 (high-recall), 64 (max).
47    #[serde(default)]
48    pub hnsw_m: Option<u32>,
49    /// HNSW ef_construction — candidate pool size during build. `None` = default (150).
50    /// Higher ef_construction → better graph quality, slower build.
51    /// Recommended values: 100 (fast), 150 (default), 200 (quality), 400 (max quality).
52    #[serde(default)]
53    pub hnsw_ef_construction: Option<u32>,
54}
55
56impl VectorStoragePolicy {
57    pub fn default_f16(column: &str, dim: u32, metric: VectorMetric) -> Self {
58        Self {
59            column_name: column.to_string(),
60            dim,
61            metric,
62            precision: VectorPrecision::F16,
63            pq: None,
64            keep_raw_for_reranking: true,
65            pre_normalize: false,
66            hnsw_m: None,
67            hnsw_ef_construction: None,
68        }
69    }
70}
71
72/// Product Quantization configuration
73#[derive(Debug, Clone, Serialize, Deserialize)]
74pub struct PQConfig {
75    /// Number of sub-vectors M (dim must be divisible by M)
76    pub num_subvectors: usize,
77    /// Bits per code (8 = 256 centroids per sub-vector)
78    pub bits_per_code: u8,
79    /// Number of training samples for codebook
80    pub train_sample_size: usize,
81}
82
83/// Marker struct for documentation purposes — actual schema is enforced by
84/// column names in llm_columns module.
85pub struct LlmContextSchema;