1use crate::db_postgres::{Chunk, Document};
2use crate::embedding::EmbeddingModelType;
3use rocket::fs::TempFile;
4use rocket::serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use utoipa::ToSchema;
7
8#[derive(FromForm)]
9pub struct Upload<'r> {
10 pub collection_name: String,
11 pub file: TempFile<'r>,
12 pub metadata: Option<HashMap<String, String>>,
13 pub model: Option<String>,
14}
15
16#[derive(Debug, Clone, Serialize, Deserialize)]
17#[serde(crate = "rocket::serde")]
18#[derive(ToSchema)]
19pub struct SimilaritySearchDocsResponse {
20 pub results: Vec<(f64, Document)>,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize)]
25#[serde(crate = "rocket::serde")]
26#[derive(ToSchema)]
27pub struct ListResponse {
28 pub results: Vec<Document>,
30}
31
32#[derive(Debug, Clone, Serialize, Deserialize)]
33#[serde(crate = "rocket::serde")]
34#[derive(ToSchema)]
35pub struct ExecuteSqlRequest {
36 pub sql_query: String,
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
40#[serde(crate = "rocket::serde")]
41#[derive(ToSchema)]
42pub struct ExecuteSqlResponse {
43 pub results: Vec<serde_json::Value>,
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
47#[serde(crate = "rocket::serde")]
48#[derive(ToSchema)]
49pub struct GetInfoRequest {
50 pub collection_name: String,
51}
52
53#[derive(Debug, Clone, Serialize, Deserialize)]
54#[serde(crate = "rocket::serde")]
55#[derive(ToSchema)]
56pub struct GetInfoResponse {
57 pub table_name: String,
58 pub dimension: usize,
59 pub schema: HashMap<String, String>,
60 pub num_documents: usize,
61}
62
63#[derive(Debug, Clone, Serialize, Deserialize)]
64#[serde(crate = "rocket::serde")]
65#[derive(ToSchema)]
66pub struct AddRequest {
67 pub collection_name: String,
68 pub texts: Vec<String>,
69 pub vectors: Option<Vec<Vec<f32>>>,
70 pub metadata: Option<Vec<HashMap<String, String>>>,
71 pub model: Option<EmbeddingModelType>,
72 pub document_id: Option<String>,
73 pub create_new_doc: Option<bool>,
74}
75
76#[derive(Debug, Clone, Serialize, Deserialize)]
77#[serde(crate = "rocket::serde")]
78#[derive(ToSchema)]
79pub struct AddResponse {
80 pub id: String,
81}
82
83#[derive(Debug, Clone, Serialize, Deserialize)]
84#[serde(crate = "rocket::serde")]
85#[derive(ToSchema)]
86pub struct UpdateRequest {
87 pub collection_name: String,
88 pub id: String,
89 pub text: Option<String>,
90 pub metadata: Option<HashMap<String, String>>,
91 pub model: Option<EmbeddingModelType>,
92}
93
94#[derive(Debug, Clone, Serialize, Deserialize)]
95#[serde(crate = "rocket::serde")]
96#[derive(ToSchema)]
97pub struct DeleteRequest {
98 pub collection_name: String,
99 pub ids: Vec<String>,
100}
101
102#[derive(Debug, Clone, Serialize, Deserialize)]
103#[serde(crate = "rocket::serde")]
104#[derive(ToSchema)]
105pub struct DeleteResponse {
106 pub deleted_count: usize,
107}
108
109#[derive(Debug, Clone, Serialize, Deserialize)]
110#[serde(crate = "rocket::serde")]
111#[derive(ToSchema)]
112pub struct UpdateResponse {
113 pub success: bool,
114}
115
116#[derive(Debug, Clone, Serialize, Deserialize)]
117#[serde(crate = "rocket::serde")]
118#[derive(ToSchema)]
119pub struct GetDocsRequest {
120 pub collection_name: String,
121 pub ids: Option<Vec<String>>,
122 pub limit: Option<usize>,
123 pub offset: Option<usize>,
124}
125
126#[derive(Debug, Clone, Serialize, Deserialize)]
127#[serde(crate = "rocket::serde")]
128#[derive(ToSchema)]
129pub struct GetDocsResponse {
130 pub results: Vec<Document>,
131}
132
133#[derive(Debug, Clone, Serialize, Deserialize)]
134#[serde(crate = "rocket::serde")]
135#[derive(ToSchema)]
136pub struct GetTextsRequest {
137 pub collection_name: String,
138 pub ids: Option<Vec<String>>,
139 pub limit: Option<usize>,
140 pub offset: Option<usize>,
141}
142
143#[derive(Debug, Clone, Serialize, Deserialize)]
144#[serde(crate = "rocket::serde")]
145#[derive(ToSchema)]
146pub struct GetTextsResponse {
147 pub results: Vec<Chunk>,
148}
149
150#[derive(Debug, Clone, Serialize, Deserialize)]
151#[serde(crate = "rocket::serde")]
152#[derive(ToSchema)]
153pub struct CountItemsRequest {
154 pub collection_name: String,
155}
156
157#[derive(Debug, Clone, Serialize, Deserialize)]
158#[serde(crate = "rocket::serde")]
159#[derive(ToSchema)]
160pub struct CountResponse {
161 pub count: usize,
162}
163
164#[derive(Debug, Clone, Serialize, Deserialize)]
165#[serde(crate = "rocket::serde")]
166#[derive(ToSchema)]
167pub struct ClearResponse {
168 pub success: bool,
169}
170
171#[derive(Debug, Clone, Serialize, Deserialize)]
172#[serde(crate = "rocket::serde")]
173#[derive(ToSchema)]
174pub struct InfoResponse {
175 pub info: HashMap<String, String>,
176}
177
178#[derive(FromForm)]
179pub struct BatchUpload<'r> {
180 pub collection_name: String,
181 pub files: Vec<TempFile<'r>>,
182 pub metadata: Option<String>,
183 pub model: Option<String>,
184}
185
186#[derive(Debug, Clone, Serialize, Deserialize)]
187#[serde(crate = "rocket::serde")]
188#[derive(ToSchema)]
189pub struct BatchAddFileResponse {
190 pub ids: Vec<String>,
191}
192
193#[derive(Debug, Clone, Serialize, Deserialize)]
194#[serde(crate = "rocket::serde")]
195#[derive(ToSchema)]
196pub enum SimilarityMethod {
197 Cosine,
198 Euclidean,
199 DotProduct,
200}
201
202#[derive(Debug, Clone, Serialize, Deserialize)]
203#[serde(crate = "rocket::serde")]
204#[derive(ToSchema)]
205pub enum SearchQuery {
206 VSS(Vec<f32>),
207 FTS(String),
208}
209
210#[derive(Debug, Copy, Clone, Serialize, Deserialize)]
211#[serde(crate = "rocket::serde")]
212#[derive(ToSchema, Default)]
213pub enum SearchType {
214 #[default]
215 Chunks,
216 Documents,
217}
218
219#[derive(Debug, Copy, Clone, Serialize, Deserialize)]
220#[serde(crate = "rocket::serde")]
221#[derive(ToSchema, Default)]
222pub enum SearchMode {
223 Vector,
224 #[default]
225 Hybrid,
226 FullText,
227}
228
229#[derive(Debug, Clone, Serialize, Deserialize)]
230#[serde(crate = "rocket::serde")]
231#[derive(ToSchema)]
232pub struct SearchRequest {
233 pub collection_name: String,
234 pub query: String,
235 pub top_k: Option<usize>,
236 pub model: Option<EmbeddingModelType>,
237 pub model_url: Option<String>,
238 pub search_type: Option<SearchType>,
239 pub search_mode: Option<SearchMode>,
240 pub similarity_method: Option<SimilarityMethod>,
241 pub document_ids: Option<Vec<String>>,
242}
243
244#[derive(Debug, Clone, Serialize, Deserialize)]
245#[serde(crate = "rocket::serde")]
246#[derive(ToSchema)]
247pub struct SearchResponse {
248 pub results: Vec<SearchResult>,
249}
250
251#[derive(Debug, Clone, Serialize, Deserialize)]
252#[serde(crate = "rocket::serde")]
253#[derive(ToSchema)]
254pub struct SearchResult {
255 pub distance: f64,
256 pub content: String,
257 pub document: Document,
258}
259
260#[derive(Debug, Clone, Serialize, Deserialize)]
261#[serde(crate = "rocket::serde")]
262#[derive(ToSchema)]
263pub struct CreateCollectionRequest {
264 pub collection_name: String,
265 pub dimensions: usize,
266}
267
268#[derive(Debug, Clone, Serialize, Deserialize)]
269#[serde(crate = "rocket::serde")]
270#[derive(ToSchema)]
271pub struct CreateCollectionResponse {
272 pub success: bool,
273}
274
275#[derive(Debug, Clone, Serialize, Deserialize)]
276#[serde(crate = "rocket::serde")]
277#[derive(ToSchema)]
278pub struct DeleteCollectionRequest {
279 pub collection_name: String,
280}
281
282#[derive(Debug, Clone, Serialize, Deserialize)]
283#[serde(crate = "rocket::serde")]
284#[derive(ToSchema)]
285pub struct DeleteCollectionResponse {
286 pub success: bool,
287}
288
289#[derive(Debug, Clone, Serialize, Deserialize)]
290#[serde(crate = "rocket::serde")]
291#[derive(ToSchema)]
292pub struct ClearCollectionRequest {
293 pub collection_name: String,
294}
295
296#[derive(Debug, Clone, Serialize, Deserialize)]
297#[serde(crate = "rocket::serde")]
298#[derive(ToSchema)]
299pub struct ClearCollectionResponse {
300 pub success: bool,
301}