Skip to main content

qwencode_rs/query/
builder.rs

1use crate::query::session::QueryHandle;
2use crate::types::config::QueryOptions;
3use anyhow::Result;
4
5/// Query builder for fluent API
6pub 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    /// Set the prompt for the query
20    pub fn prompt(mut self, prompt: &str) -> Self {
21        self.prompt = Some(prompt.to_string());
22        self
23    }
24
25    /// Set the working directory
26    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    /// Set the AI model
32    pub fn model(mut self, model: &str) -> Self {
33        self.options.model = Some(model.to_string());
34        self
35    }
36
37    /// Set the permission mode
38    pub fn permission_mode(mut self, mode: crate::types::permission::PermissionMode) -> Self {
39        self.options.permission_mode = mode;
40        self
41    }
42
43    /// Enable debug mode
44    pub fn debug(mut self, debug: bool) -> Self {
45        self.options.debug = debug;
46        self
47    }
48
49    /// Set maximum session turns
50    pub fn max_turns(mut self, max: i32) -> Self {
51        self.options.max_session_turns = max;
52        self
53    }
54
55    /// Set the session ID
56    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    /// Build the query options
62    pub fn build(self) -> Result<(Option<String>, QueryOptions)> {
63        Ok((self.prompt, self.options))
64    }
65
66    /// Execute the query (placeholder for now)
67    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        // For now, just create a session handle
75        // In the full implementation, this would spawn the CLI process
76        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}