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(&[&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}