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    #[serde(default)]
162    pub request_id: Option<String>,
163}
164
165// ── xAI Collection Proxy Types ──────────────────────────────────────────────
166
167/// A user-scoped xAI collection (proxied through quantum-ai).
168#[derive(Debug, Clone, Serialize, Deserialize)]
169pub struct Collection {
170    /// Collection ID (xAI-issued).
171    pub id: String,
172
173    /// Human-readable name.
174    pub name: String,
175
176    /// Optional description.
177    #[serde(default)]
178    pub description: Option<String>,
179
180    /// Number of documents in the collection.
181    #[serde(default)]
182    pub document_count: Option<u64>,
183
184    /// Owner: user ID or "shared".
185    #[serde(default)]
186    pub owner: Option<String>,
187
188    /// Backend provider (e.g. "xai").
189    #[serde(default)]
190    pub provider: Option<String>,
191
192    /// ISO timestamp.
193    #[serde(default)]
194    pub created_at: Option<String>,
195}
196
197/// Request body for creating a collection.
198#[derive(Debug, Clone, Serialize)]
199pub struct CreateCollectionRequest {
200    pub name: String,
201}
202
203/// A document within a collection.
204#[derive(Debug, Clone, Serialize, Deserialize)]
205pub struct CollectionDocument {
206    pub file_id: String,
207    pub name: String,
208    #[serde(default)]
209    pub size_bytes: Option<u64>,
210    #[serde(default)]
211    pub content_type: Option<String>,
212    #[serde(default)]
213    pub processing_status: Option<String>,
214    #[serde(default)]
215    pub document_status: Option<String>,
216    #[serde(default)]
217    pub indexed: Option<bool>,
218    #[serde(default)]
219    pub created_at: Option<String>,
220}
221
222/// A search result from collection search.
223#[derive(Debug, Clone, Deserialize)]
224pub struct CollectionSearchResult {
225    pub content: String,
226    #[serde(default)]
227    pub score: Option<f64>,
228    #[serde(default)]
229    pub file_id: Option<String>,
230    #[serde(default)]
231    pub collection_id: Option<String>,
232    #[serde(default)]
233    pub metadata: Option<serde_json::Value>,
234}
235
236/// Request body for collection search.
237#[derive(Debug, Clone, Serialize)]
238pub struct CollectionSearchRequest {
239    pub query: String,
240    pub collection_ids: Vec<String>,
241    #[serde(skip_serializing_if = "Option::is_none")]
242    pub mode: Option<String>,
243    #[serde(skip_serializing_if = "Option::is_none")]
244    pub max_results: Option<usize>,
245}
246
247/// Upload result for a document added to a collection.
248#[derive(Debug, Clone, Deserialize)]
249pub struct CollectionUploadResult {
250    pub file_id: String,
251    pub filename: String,
252    #[serde(default)]
253    pub bytes: Option<u64>,
254}
255
256// Wrapper types for API responses.
257
258#[derive(Deserialize)]
259struct CollectionsListResponse {
260    collections: Vec<Collection>,
261}
262
263#[derive(Deserialize)]
264struct CollectionDocumentsResponse {
265    documents: Vec<CollectionDocument>,
266}
267
268#[derive(Deserialize)]
269struct CollectionSearchResponse {
270    results: Vec<CollectionSearchResult>,
271}
272
273#[derive(Deserialize)]
274struct DeleteCollectionResponse {
275    #[serde(default)]
276    message: String,
277}
278
279impl Client {
280    /// Searches Vertex AI RAG corpora for relevant documentation.
281    pub async fn rag_search(&self, req: &RagSearchRequest) -> Result<RagSearchResponse> {
282        let (mut resp, meta) = self
283            .post_json::<RagSearchRequest, RagSearchResponse>("/qai/v1/rag/search", req)
284            .await?;
285        if resp.cost_ticks == 0 {
286            resp.cost_ticks = meta.cost_ticks;
287        }
288        if resp.request_id.is_empty() {
289            resp.request_id = meta.request_id;
290        }
291        Ok(resp)
292    }
293
294    /// Lists available Vertex AI RAG corpora.
295    pub async fn rag_corpora(&self) -> Result<Vec<RagCorpus>> {
296        let (resp, _meta) = self
297            .get_json::<RagCorporaResponse>("/qai/v1/rag/corpora")
298            .await?;
299        Ok(resp.corpora)
300    }
301
302    /// Searches provider API documentation via SurrealDB vector search.
303    pub async fn surreal_rag_search(
304        &self,
305        req: &SurrealRagSearchRequest,
306    ) -> Result<SurrealRagSearchResponse> {
307        let (mut resp, meta) = self
308            .post_json::<SurrealRagSearchRequest, SurrealRagSearchResponse>(
309                "/qai/v1/rag/surreal/search",
310                req,
311            )
312            .await?;
313        if resp.cost_ticks == 0 {
314            resp.cost_ticks = meta.cost_ticks;
315        }
316        if resp.request_id.is_empty() {
317            resp.request_id = meta.request_id;
318        }
319        Ok(resp)
320    }
321
322    /// Lists available SurrealDB RAG documentation providers.
323    pub async fn surreal_rag_providers(&self) -> Result<SurrealRagProvidersResponse> {
324        let (resp, _meta) = self
325            .get_json::<SurrealRagProvidersResponse>("/qai/v1/rag/surreal/providers")
326            .await?;
327        Ok(resp)
328    }
329
330    // ── xAI Collection Proxy (user-scoped) ──────────────────────────────────
331
332    /// Lists the user's collections plus shared collections.
333    pub async fn collections_list(&self) -> Result<Vec<Collection>> {
334        let (resp, _meta) = self
335            .get_json::<CollectionsListResponse>("/qai/v1/rag/collections")
336            .await?;
337        Ok(resp.collections)
338    }
339
340    /// Creates a new user-owned collection.
341    pub async fn collections_create(&self, name: &str) -> Result<Collection> {
342        let req = CreateCollectionRequest {
343            name: name.to_string(),
344        };
345        let (resp, _meta) = self
346            .post_json::<CreateCollectionRequest, Collection>("/qai/v1/rag/collections", &req)
347            .await?;
348        Ok(resp)
349    }
350
351    /// Gets details for a single collection (must be owned or shared).
352    pub async fn collections_get(&self, id: &str) -> Result<Collection> {
353        let (resp, _meta) = self
354            .get_json::<Collection>(&format!("/qai/v1/rag/collections/{id}"))
355            .await?;
356        Ok(resp)
357    }
358
359    /// Deletes a collection (owner only).
360    pub async fn collections_delete(&self, id: &str) -> Result<String> {
361        let (resp, _meta) = self
362            .delete_json::<DeleteCollectionResponse>(&format!("/qai/v1/rag/collections/{id}"))
363            .await?;
364        Ok(resp.message)
365    }
366
367    /// Lists documents in a collection.
368    pub async fn collections_documents(&self, collection_id: &str) -> Result<Vec<CollectionDocument>> {
369        let (resp, _meta) = self
370            .get_json::<CollectionDocumentsResponse>(&format!(
371                "/qai/v1/rag/collections/{collection_id}/documents"
372            ))
373            .await?;
374        Ok(resp.documents)
375    }
376
377    /// Uploads a file to a collection. The server handles the two-step
378    /// xAI upload (files API + management API) with the master key.
379    pub async fn collections_upload(
380        &self,
381        collection_id: &str,
382        filename: &str,
383        content: Vec<u8>,
384    ) -> Result<CollectionUploadResult> {
385        let part = reqwest::multipart::Part::bytes(content)
386            .file_name(filename.to_string())
387            .mime_str("application/octet-stream")
388            .map_err(|e| crate::error::Error::Api(crate::error::ApiError {
389                status_code: 0,
390                code: "multipart_error".into(),
391                message: e.to_string(),
392                request_id: String::new(),
393            }))?;
394        let form = reqwest::multipart::Form::new().part("file", part);
395        let (resp, _meta) = self
396            .post_multipart::<CollectionUploadResult>(
397                &format!("/qai/v1/rag/collections/{collection_id}/upload"),
398                form,
399            )
400            .await?;
401        Ok(resp)
402    }
403
404    /// Searches across collections (user's + shared) with hybrid/semantic/keyword mode.
405    pub async fn collections_search(
406        &self,
407        req: &CollectionSearchRequest,
408    ) -> Result<Vec<CollectionSearchResult>> {
409        let (resp, _meta) = self
410            .post_json::<CollectionSearchRequest, CollectionSearchResponse>(
411                "/qai/v1/rag/search/collections",
412                req,
413            )
414            .await?;
415        Ok(resp.results)
416    }
417}
quantum_sdk/rag.rs

quantum_sdk/
rag.rs
