qwencode_rs/query/
builder.rs1use crate::query::session::QueryHandle;
2use crate::types::config::QueryOptions;
3use anyhow::Result;
4
5pub struct QueryBuilder {
7 options: QueryOptions,
8 pub prompt: Option<String>,
9}
10
11impl QueryBuilder {
12 pub fn new() -> Self {
13 QueryBuilder {
14 options: QueryOptions::default(),
15 prompt: None,
16 }
17 }
18
19 pub fn prompt(mut self, prompt: &str) -> Self {
21 self.prompt = Some(prompt.to_string());
22 self
23 }
24
25 pub fn cwd<P: Into<std::path::PathBuf>>(mut self, path: P) -> Self {
27 self.options.cwd = Some(path.into());
28 self
29 }
30
31 pub fn model(mut self, model: &str) -> Self {
33 self.options.model = Some(model.to_string());
34 self
35 }
36
37 pub fn permission_mode(mut self, mode: crate::types::permission::PermissionMode) -> Self {
39 self.options.permission_mode = mode;
40 self
41 }
42
43 pub fn debug(mut self, debug: bool) -> Self {
45 self.options.debug = debug;
46 self
47 }
48
49 pub fn max_turns(mut self, max: i32) -> Self {
51 self.options.max_session_turns = max;
52 self
53 }
54
55 pub fn session_id(mut self, session_id: &str) -> Self {
57 self.options.session_id = Some(session_id.to_string());
58 self
59 }
60
61 pub fn build(self) -> Result<(Option<String>, QueryOptions)> {
63 Ok((self.prompt, self.options))
64 }
65
66 pub async fn execute(self) -> Result<QueryHandle> {
68 let (prompt, options) = self.build()?;
69
70 if let Some(ref p) = prompt {
71 tracing::info!("Executing query with prompt: {}", p);
72 }
73
74 let handle = QueryHandle::new(options.session_id.clone());
77 Ok(handle)
78 }
79}
80
81impl Default for QueryBuilder {
82 fn default() -> Self {
83 QueryBuilder::new()
84 }
85}
86
87#[cfg(test)]
88mod tests {
89 use super::*;
90 use crate::types::permission::PermissionMode;
91
92 #[test]
93 fn test_query_builder_default() {
94 let builder = QueryBuilder::new();
95 assert!(builder.prompt.is_none());
96 }
97
98 #[test]
99 fn test_query_builder_prompt() {
100 let builder = QueryBuilder::new().prompt("Hello, world!");
101
102 assert_eq!(builder.prompt, Some("Hello, world!".to_string()));
103 }
104
105 #[test]
106 fn test_query_builder_cwd() {
107 let builder = QueryBuilder::new().cwd("/tmp/test");
108
109 assert_eq!(
110 builder.options.cwd,
111 Some(std::path::PathBuf::from("/tmp/test"))
112 );
113 }
114
115 #[test]
116 fn test_query_builder_model() {
117 let builder = QueryBuilder::new().model("qwen-max");
118
119 assert_eq!(builder.options.model, Some("qwen-max".to_string()));
120 }
121
122 #[test]
123 fn test_query_builder_permission_mode() {
124 let builder = QueryBuilder::new().permission_mode(PermissionMode::Yolo);
125
126 assert_eq!(builder.options.permission_mode, PermissionMode::Yolo);
127 }
128
129 #[test]
130 fn test_query_builder_debug() {
131 let builder = QueryBuilder::new().debug(true);
132
133 assert!(builder.options.debug);
134 }
135
136 #[test]
137 fn test_query_builder_max_turns() {
138 let builder = QueryBuilder::new().max_turns(10);
139
140 assert_eq!(builder.options.max_session_turns, 10);
141 }
142
143 #[test]
144 fn test_query_builder_session_id() {
145 let builder = QueryBuilder::new().session_id("my-session");
146
147 assert_eq!(builder.options.session_id, Some("my-session".to_string()));
148 }
149
150 #[test]
151 fn test_query_builder_build() {
152 let (prompt, options) = QueryBuilder::new()
153 .prompt("Test")
154 .model("qwen-plus")
155 .debug(true)
156 .build()
157 .unwrap();
158
159 assert_eq!(prompt, Some("Test".to_string()));
160 assert_eq!(options.model, Some("qwen-plus".to_string()));
161 assert!(options.debug);
162 }
163
164 #[tokio::test]
165 async fn test_query_builder_execute() {
166 let handle = QueryBuilder::new()
167 .prompt("Test query")
168 .session_id("test-session")
169 .execute()
170 .await
171 .unwrap();
172
173 assert_eq!(handle.session_id(), "test-session");
174 assert!(!handle.is_closed());
175 }
176
177 #[test]
178 fn test_query_builder_chained_calls() {
179 let builder = QueryBuilder::new()
180 .prompt("Hello")
181 .model("qwen-max")
182 .cwd("/tmp")
183 .permission_mode(PermissionMode::AutoEdit)
184 .debug(true)
185 .max_turns(5)
186 .session_id("chained-test");
187
188 assert_eq!(builder.prompt, Some("Hello".to_string()));
189 assert_eq!(builder.options.model, Some("qwen-max".to_string()));
190 assert_eq!(builder.options.cwd, Some(std::path::PathBuf::from("/tmp")));
191 assert_eq!(builder.options.permission_mode, PermissionMode::AutoEdit);
192 assert!(builder.options.debug);
193 assert_eq!(builder.options.max_session_turns, 5);
194 assert_eq!(builder.options.session_id, Some("chained-test".to_string()));
195 }
196}