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 SurrealRagProvider {
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/// Response from listing SurrealDB RAG providers.
155#[derive(Debug, Clone, Deserialize)]
156pub struct SurrealRagProvidersResponse {
157    pub providers: Vec<SurrealRagProvider>,
158}
159
160impl Client {
161    /// Searches Vertex AI RAG corpora for relevant documentation.
162    pub async fn rag_search(&self, req: &RagSearchRequest) -> Result<RagSearchResponse> {
163        let (mut resp, meta) = self
164            .post_json::<RagSearchRequest, RagSearchResponse>("/qai/v1/rag/search", req)
165            .await?;
166        if resp.cost_ticks == 0 {
167            resp.cost_ticks = meta.cost_ticks;
168        }
169        if resp.request_id.is_empty() {
170            resp.request_id = meta.request_id;
171        }
172        Ok(resp)
173    }
174
175    /// Lists available Vertex AI RAG corpora.
176    pub async fn rag_corpora(&self) -> Result<Vec<RagCorpus>> {
177        let (resp, _meta) = self
178            .get_json::<RagCorporaResponse>("/qai/v1/rag/corpora")
179            .await?;
180        Ok(resp.corpora)
181    }
182
183    /// Searches provider API documentation via SurrealDB vector search.
184    pub async fn surreal_rag_search(
185        &self,
186        req: &SurrealRagSearchRequest,
187    ) -> Result<SurrealRagSearchResponse> {
188        let (mut resp, meta) = self
189            .post_json::<SurrealRagSearchRequest, SurrealRagSearchResponse>(
190                "/qai/v1/rag/surreal/search",
191                req,
192            )
193            .await?;
194        if resp.cost_ticks == 0 {
195            resp.cost_ticks = meta.cost_ticks;
196        }
197        if resp.request_id.is_empty() {
198            resp.request_id = meta.request_id;
199        }
200        Ok(resp)
201    }
202
203    /// Lists available SurrealDB RAG documentation providers.
204    pub async fn surreal_rag_providers(&self) -> Result<SurrealRagProvidersResponse> {
205        let (resp, _meta) = self
206            .get_json::<SurrealRagProvidersResponse>("/qai/v1/rag/surreal/providers")
207            .await?;
208        Ok(resp)
209    }
210}
quantum_sdk/rag.rs

quantum_sdk/
rag.rs
