quantum_sdk/
rag.rs

1use serde::{Deserialize, Serialize};
2
3use crate::client::Client;
4use crate::error::Result;
5
6/// Request body for Vertex AI RAG search.
7#[derive(Debug, Clone, Serialize, Default)]
8pub struct RagSearchRequest {
9    /// Search query.
10    pub query: String,
11
12    /// Filter by corpus name or ID (fuzzy match). Omit to search all corpora.
13    #[serde(skip_serializing_if = "Option::is_none")]
14    pub corpus: Option<String>,
15
16    /// Maximum number of results to return (default 10).
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub top_k: Option<i32>,
19}
20
21/// Response from RAG search.
22#[derive(Debug, Clone, Deserialize)]
23pub struct RagSearchResponse {
24    /// Matching document chunks.
25    pub results: Vec<RagResult>,
26
27    /// Original search query.
28    pub query: String,
29
30    /// Corpora that were searched.
31    #[serde(default)]
32    pub corpora: Option<Vec<String>>,
33
34    /// Total cost in ticks.
35    #[serde(default)]
36    pub cost_ticks: i64,
37
38    /// Unique request identifier.
39    #[serde(default)]
40    pub request_id: String,
41}
42
43/// A single result from RAG search.
44#[derive(Debug, Clone, Deserialize)]
45pub struct RagResult {
46    /// Source document URI.
47    pub source_uri: String,
48
49    /// Display name of the source.
50    pub source_name: String,
51
52    /// Matching text chunk.
53    pub text: String,
54
55    /// Relevance score.
56    pub score: f64,
57
58    /// Vector distance (lower is more similar).
59    pub distance: f64,
60}
61
62/// Describes an available RAG corpus.
63#[derive(Debug, Clone, Deserialize)]
64pub struct RagCorpus {
65    /// Full resource name.
66    pub name: String,
67
68    /// Human-readable name.
69    #[serde(rename = "displayName")]
70    pub display_name: String,
71
72    /// Describes the corpus contents.
73    pub description: String,
74
75    /// Corpus state (e.g. "ACTIVE").
76    pub state: String,
77}
78
79#[derive(Deserialize)]
80struct RagCorporaResponse {
81    corpora: Vec<RagCorpus>,
82}
83
84/// Request body for SurrealDB-backed RAG search.
85#[derive(Debug, Clone, Serialize, Default)]
86pub struct SurrealRagSearchRequest {
87    /// Search query.
88    pub query: String,
89
90    /// Filter by documentation provider (e.g. "xai", "claude", "heygen").
91    #[serde(skip_serializing_if = "Option::is_none")]
92    pub provider: Option<String>,
93
94    /// Maximum number of results (default 10, max 50).
95    #[serde(skip_serializing_if = "Option::is_none")]
96    pub limit: Option<i32>,
97}
98
99/// Response from SurrealDB RAG search.
100#[derive(Debug, Clone, Deserialize)]
101pub struct SurrealRagSearchResponse {
102    /// Matching documentation chunks.
103    pub results: Vec<SurrealRagResult>,
104
105    /// Original search query.
106    pub query: String,
107
108    /// Provider filter that was applied.
109    #[serde(default)]
110    pub provider: Option<String>,
111
112    /// Total cost in ticks.
113    #[serde(default)]
114    pub cost_ticks: i64,
115
116    /// Unique request identifier.
117    #[serde(default)]
118    pub request_id: String,
119}
120
121/// A single result from SurrealDB RAG search.
122#[derive(Debug, Clone, Deserialize)]
123pub struct SurrealRagResult {
124    /// Documentation provider.
125    pub provider: String,
126
127    /// Document title.
128    pub title: String,
129
130    /// Section heading.
131    pub heading: String,
132
133    /// Original source file path.
134    pub source_file: String,
135
136    /// Matching text chunk.
137    pub content: String,
138
139    /// Cosine similarity score.
140    pub score: f64,
141}
142
143/// A SurrealDB RAG provider.
144#[derive(Debug, Clone, Deserialize)]
145pub struct SurrealRagProviderInfo {
146    /// Provider identifier (e.g. "xai", "claude").
147    pub provider: String,
148
149    /// Number of document chunks for this provider.
150    #[serde(default)]
151    pub chunk_count: Option<i64>,
152}
153
154/// Backwards-compatible alias.
155pub type SurrealRagProvider = SurrealRagProviderInfo;
156
157/// Response from listing SurrealDB RAG providers.
158#[derive(Debug, Clone, Deserialize)]
159pub struct SurrealRagProvidersResponse {
160    pub providers: Vec<SurrealRagProviderInfo>,
161}
162
163// ── xAI Collection Proxy Types ──────────────────────────────────────────────
164
165/// A user-scoped xAI collection (proxied through quantum-ai).
166#[derive(Debug, Clone, Serialize, Deserialize)]
167pub struct Collection {
168    /// Collection ID (xAI-issued).
169    pub id: String,
170
171    /// Human-readable name.
172    pub name: String,
173
174    /// Optional description.
175    #[serde(default)]
176    pub description: Option<String>,
177
178    /// Number of documents in the collection.
179    #[serde(default)]
180    pub document_count: Option<u64>,
181
182    /// Owner: user ID or "shared".
183    #[serde(default)]
184    pub owner: Option<String>,
185
186    /// Backend provider (e.g. "xai").
187    #[serde(default)]
188    pub provider: Option<String>,
189
190    /// ISO timestamp.
191    #[serde(default)]
192    pub created_at: Option<String>,
193}
194
195/// Request body for creating a collection.
196#[derive(Debug, Clone, Serialize)]
197pub struct CreateCollectionRequest {
198    pub name: String,
199}
200
201/// A document within a collection.
202#[derive(Debug, Clone, Serialize, Deserialize)]
203pub struct CollectionDocument {
204    pub file_id: String,
205    pub name: String,
206    #[serde(default)]
207    pub size_bytes: Option<u64>,
208    #[serde(default)]
209    pub content_type: Option<String>,
210    #[serde(default)]
211    pub processing_status: Option<String>,
212    #[serde(default)]
213    pub document_status: Option<String>,
214    #[serde(default)]
215    pub indexed: Option<bool>,
216    #[serde(default)]
217    pub created_at: Option<String>,
218}
219
220/// A search result from collection search.
221#[derive(Debug, Clone, Deserialize)]
222pub struct CollectionSearchResult {
223    pub content: String,
224    #[serde(default)]
225    pub score: Option<f64>,
226    #[serde(default)]
227    pub file_id: Option<String>,
228    #[serde(default)]
229    pub collection_id: Option<String>,
230    #[serde(default)]
231    pub metadata: Option<serde_json::Value>,
232}
233
234/// Request body for collection search.
235#[derive(Debug, Clone, Serialize)]
236pub struct CollectionSearchRequest {
237    pub query: String,
238    pub collection_ids: Vec<String>,
239    #[serde(skip_serializing_if = "Option::is_none")]
240    pub mode: Option<String>,
241    #[serde(skip_serializing_if = "Option::is_none")]
242    pub max_results: Option<usize>,
243}
244
245/// Upload result for a document added to a collection.
246#[derive(Debug, Clone, Deserialize)]
247pub struct CollectionUploadResult {
248    pub file_id: String,
249    pub filename: String,
250    #[serde(default)]
251    pub bytes: Option<u64>,
252}
253
254// Wrapper types for API responses.
255
256#[derive(Deserialize)]
257struct CollectionsListResponse {
258    collections: Vec<Collection>,
259}
260
261#[derive(Deserialize)]
262struct CollectionDocumentsResponse {
263    documents: Vec<CollectionDocument>,
264}
265
266#[derive(Deserialize)]
267struct CollectionSearchResponse {
268    results: Vec<CollectionSearchResult>,
269}
270
271#[derive(Deserialize)]
272struct DeleteCollectionResponse {
273    #[serde(default)]
274    message: String,
275}
276
277impl Client {
278    /// Searches Vertex AI RAG corpora for relevant documentation.
279    pub async fn rag_search(&self, req: &RagSearchRequest) -> Result<RagSearchResponse> {
280        let (mut resp, meta) = self
281            .post_json::<RagSearchRequest, RagSearchResponse>("/qai/v1/rag/search", req)
282            .await?;
283        if resp.cost_ticks == 0 {
284            resp.cost_ticks = meta.cost_ticks;
285        }
286        if resp.request_id.is_empty() {
287            resp.request_id = meta.request_id;
288        }
289        Ok(resp)
290    }
291
292    /// Lists available Vertex AI RAG corpora.
293    pub async fn rag_corpora(&self) -> Result<Vec<RagCorpus>> {
294        let (resp, _meta) = self
295            .get_json::<RagCorporaResponse>("/qai/v1/rag/corpora")
296            .await?;
297        Ok(resp.corpora)
298    }
299
300    /// Searches provider API documentation via SurrealDB vector search.
301    pub async fn surreal_rag_search(
302        &self,
303        req: &SurrealRagSearchRequest,
304    ) -> Result<SurrealRagSearchResponse> {
305        let (mut resp, meta) = self
306            .post_json::<SurrealRagSearchRequest, SurrealRagSearchResponse>(
307                "/qai/v1/rag/surreal/search",
308                req,
309            )
310            .await?;
311        if resp.cost_ticks == 0 {
312            resp.cost_ticks = meta.cost_ticks;
313        }
314        if resp.request_id.is_empty() {
315            resp.request_id = meta.request_id;
316        }
317        Ok(resp)
318    }
319
320    /// Lists available SurrealDB RAG documentation providers.
321    pub async fn surreal_rag_providers(&self) -> Result<SurrealRagProvidersResponse> {
322        let (resp, _meta) = self
323            .get_json::<SurrealRagProvidersResponse>("/qai/v1/rag/surreal/providers")
324            .await?;
325        Ok(resp)
326    }
327
328    // ── xAI Collection Proxy (user-scoped) ──────────────────────────────────
329
330    /// Lists the user's collections plus shared collections.
331    pub async fn collections_list(&self) -> Result<Vec<Collection>> {
332        let (resp, _meta) = self
333            .get_json::<CollectionsListResponse>("/qai/v1/rag/collections")
334            .await?;
335        Ok(resp.collections)
336    }
337
338    /// Creates a new user-owned collection.
339    pub async fn collections_create(&self, name: &str) -> Result<Collection> {
340        let req = CreateCollectionRequest {
341            name: name.to_string(),
342        };
343        let (resp, _meta) = self
344            .post_json::<CreateCollectionRequest, Collection>("/qai/v1/rag/collections", &req)
345            .await?;
346        Ok(resp)
347    }
348
349    /// Gets details for a single collection (must be owned or shared).
350    pub async fn collections_get(&self, id: &str) -> Result<Collection> {
351        let (resp, _meta) = self
352            .get_json::<Collection>(&format!("/qai/v1/rag/collections/{id}"))
353            .await?;
354        Ok(resp)
355    }
356
357    /// Deletes a collection (owner only).
358    pub async fn collections_delete(&self, id: &str) -> Result<String> {
359        let (resp, _meta) = self
360            .delete_json::<DeleteCollectionResponse>(&format!("/qai/v1/rag/collections/{id}"))
361            .await?;
362        Ok(resp.message)
363    }
364
365    /// Lists documents in a collection.
366    pub async fn collections_documents(&self, collection_id: &str) -> Result<Vec<CollectionDocument>> {
367        let (resp, _meta) = self
368            .get_json::<CollectionDocumentsResponse>(&format!(
369                "/qai/v1/rag/collections/{collection_id}/documents"
370            ))
371            .await?;
372        Ok(resp.documents)
373    }
374
375    /// Uploads a file to a collection. The server handles the two-step
376    /// xAI upload (files API + management API) with the master key.
377    pub async fn collections_upload(
378        &self,
379        collection_id: &str,
380        filename: &str,
381        content: Vec<u8>,
382    ) -> Result<CollectionUploadResult> {
383        let part = reqwest::multipart::Part::bytes(content)
384            .file_name(filename.to_string())
385            .mime_str("application/octet-stream")
386            .map_err(|e| crate::error::Error::Api(crate::error::ApiError {
387                status_code: 0,
388                code: "multipart_error".into(),
389                message: e.to_string(),
390                request_id: String::new(),
391            }))?;
392        let form = reqwest::multipart::Form::new().part("file", part);
393        let (resp, _meta) = self
394            .post_multipart::<CollectionUploadResult>(
395                &format!("/qai/v1/rag/collections/{collection_id}/upload"),
396                form,
397            )
398            .await?;
399        Ok(resp)
400    }
401
402    /// Searches across collections (user's + shared) with hybrid/semantic/keyword mode.
403    pub async fn collections_search(
404        &self,
405        req: &CollectionSearchRequest,
406    ) -> Result<Vec<CollectionSearchResult>> {
407        let (resp, _meta) = self
408            .post_json::<CollectionSearchRequest, CollectionSearchResponse>(
409                "/qai/v1/rag/search/collections",
410                req,
411            )
412            .await?;
413        Ok(resp.results)
414    }
415}
quantum_sdk/rag.rs

quantum_sdk/
rag.rs
