daimon_plugin_opensearch/
builder.rs1use 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
11pub 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 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 pub fn index(mut self, index: impl Into<String>) -> Self {
59 self.index = index.into();
60 self
61 }
62
63 pub fn space_type(mut self, space_type: SpaceType) -> Self {
65 self.space_type = space_type;
66 self
67 }
68
69 pub fn engine(mut self, engine: Engine) -> Self {
71 self.engine = engine;
72 self
73 }
74
75 pub fn auto_create_index(mut self, enabled: bool) -> Self {
81 self.auto_create_index = enabled;
82 self
83 }
84
85 pub fn hnsw_m(mut self, m: usize) -> Self {
88 self.hnsw_m = Some(m);
89 self
90 }
91
92 pub fn hnsw_ef_construction(mut self, ef: usize) -> Self {
95 self.hnsw_ef_construction = Some(ef);
96 self
97 }
98
99 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 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}