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(&[
129                &self.index
130            ]))
131            .send()
132            .await
133            .map_err(|e| DaimonError::Other(format!("opensearch index check error: {e}")))?;
134
135        if exists.status_code().is_success() {
136            tracing::debug!("opensearch: index '{}' already exists", self.index);
137            return Ok(());
138        }
139
140        tracing::info!("opensearch: creating k-NN index '{}'", self.index);
141
142        let body = index_settings::create_index_body(
143            self.dimensions,
144            self.space_type,
145            self.engine,
146            self.hnsw_m,
147            self.hnsw_ef_construction,
148        );
149
150        let response = client
151            .indices()
152            .create(opensearch::indices::IndicesCreateParts::Index(&self.index))
153            .body(body)
154            .send()
155            .await
156            .map_err(|e| DaimonError::Other(format!("opensearch index create error: {e}")))?;
157
158        let status = response.status_code();
159        if !status.is_success() {
160            let text = response
161                .text()
162                .await
163                .unwrap_or_else(|_| "unknown error".into());
164            return Err(DaimonError::Other(format!(
165                "opensearch index creation failed ({status}): {text}"
166            )));
167        }
168
169        tracing::info!("opensearch: index '{}' created", self.index);
170        Ok(())
171    }
172}