Skip to main content

daimon_plugin_opensearch/
builder.rs

1//! Builder for [`OpenSearchVectorStore`].
2
3use daimon_core::{DaimonError, Result};
4use opensearch::OpenSearch;
5use opensearch::http::transport::Transport;
6
7use crate::index_settings;
8use crate::store::OpenSearchVectorStore;
9use crate::{Engine, SpaceType};
10
11/// Builds an [`OpenSearchVectorStore`] with optional auto-index-creation.
12///
13/// # Example
14///
15/// ```ignore
16/// use daimon_plugin_opensearch::{OpenSearchVectorStoreBuilder, SpaceType, Engine};
17///
18/// let store = OpenSearchVectorStoreBuilder::new("http://localhost:9200", 1536)
19///     .index("embeddings")
20///     .space_type(SpaceType::CosineSimilarity)
21///     .engine(Engine::Lucene)
22///     .hnsw_m(16)
23///     .hnsw_ef_construction(256)
24///     .auto_create_index(true)
25///     .build()
26///     .await?;
27/// ```
28pub struct OpenSearchVectorStoreBuilder {
29    url: String,
30    dimensions: usize,
31    index: String,
32    space_type: SpaceType,
33    engine: Engine,
34    auto_create_index: bool,
35    hnsw_m: Option<usize>,
36    hnsw_ef_construction: Option<usize>,
37}
38
39impl OpenSearchVectorStoreBuilder {
40    /// Creates a new builder.
41    ///
42    /// - `url`: OpenSearch cluster URL (e.g. `"http://localhost:9200"`)
43    /// - `dimensions`: the fixed vector dimension count (must match your embedding model)
44    pub fn new(url: impl Into<String>, dimensions: usize) -> Self {
45        Self {
46            url: url.into(),
47            dimensions,
48            index: "daimon_vectors".into(),
49            space_type: SpaceType::default(),
50            engine: Engine::default(),
51            auto_create_index: true,
52            hnsw_m: None,
53            hnsw_ef_construction: None,
54        }
55    }
56
57    /// Sets the index name. Default: `"daimon_vectors"`.
58    pub fn index(mut self, index: impl Into<String>) -> Self {
59        self.index = index.into();
60        self
61    }
62
63    /// Sets the k-NN space type (distance metric). Default: [`SpaceType::CosineSimilarity`].
64    pub fn space_type(mut self, space_type: SpaceType) -> Self {
65        self.space_type = space_type;
66        self
67    }
68
69    /// Sets the k-NN engine. Default: [`Engine::Lucene`].
70    pub fn engine(mut self, engine: Engine) -> Self {
71        self.engine = engine;
72        self
73    }
74
75    /// Enables or disables automatic index creation on first use.
76    /// Default: `true`.
77    ///
78    /// When disabled, use the JSON from [`crate::index_settings`] to create
79    /// the index manually.
80    pub fn auto_create_index(mut self, enabled: bool) -> Self {
81        self.auto_create_index = enabled;
82        self
83    }
84
85    /// Sets the HNSW `m` parameter (max connections per layer).
86    /// `None` uses the engine default.
87    pub fn hnsw_m(mut self, m: usize) -> Self {
88        self.hnsw_m = Some(m);
89        self
90    }
91
92    /// Sets the HNSW `ef_construction` parameter (build-time search width).
93    /// `None` uses the engine default.
94    pub fn hnsw_ef_construction(mut self, ef: usize) -> Self {
95        self.hnsw_ef_construction = Some(ef);
96        self
97    }
98
99    /// Builds an [`OpenSearchVectorStore`] from a pre-existing [`OpenSearch`] client.
100    ///
101    /// Use this when you need custom transport configuration (e.g. AWS SigV4,
102    /// custom certificates, connection pool tuning).
103    pub async fn build_with_client(self, client: OpenSearch) -> Result<OpenSearchVectorStore> {
104        if self.auto_create_index {
105            self.ensure_index(&client).await?;
106        }
107
108        Ok(OpenSearchVectorStore {
109            client,
110            index: self.index,
111            dimensions: self.dimensions,
112            space_type: self.space_type,
113        })
114    }
115
116    /// Builds the [`OpenSearchVectorStore`], optionally creating the index.
117    pub async fn build(self) -> Result<OpenSearchVectorStore> {
118        let transport = Transport::single_node(&self.url)
119            .map_err(|e| DaimonError::Other(format!("opensearch transport error: {e}")))?;
120        let client = OpenSearch::new(transport);
121
122        self.build_with_client(client).await
123    }
124
125    async fn ensure_index(&self, client: &OpenSearch) -> Result<()> {
126        let exists = client
127            .indices()
128            .exists(opensearch::indices::IndicesExistsParts::Index(&[&self.index]))
129            .send()
130            .await
131            .map_err(|e| DaimonError::Other(format!("opensearch index check error: {e}")))?;
132
133        if exists.status_code().is_success() {
134            tracing::debug!("opensearch: index '{}' already exists", self.index);
135            return Ok(());
136        }
137
138        tracing::info!("opensearch: creating k-NN index '{}'", self.index);
139
140        let body = index_settings::create_index_body(
141            self.dimensions,
142            self.space_type,
143            self.engine,
144            self.hnsw_m,
145            self.hnsw_ef_construction,
146        );
147
148        let response = client
149            .indices()
150            .create(opensearch::indices::IndicesCreateParts::Index(&self.index))
151            .body(body)
152            .send()
153            .await
154            .map_err(|e| DaimonError::Other(format!("opensearch index create error: {e}")))?;
155
156        let status = response.status_code();
157        if !status.is_success() {
158            let text = response
159                .text()
160                .await
161                .unwrap_or_else(|_| "unknown error".into());
162            return Err(DaimonError::Other(format!(
163                "opensearch index creation failed ({status}): {text}"
164            )));
165        }
166
167        tracing::info!("opensearch: index '{}' created", self.index);
168        Ok(())
169    }
170}