1 files changed, 41 insertions, 0 deletions
diff --git a/src/config/input.rs b/src/config/input.rs
index d6d154b..71807f3 100644
--- a/src/config/input.rs
+++ b/src/config/input.rs
@@ -27,6 +27,13 @@ pub struct Input {
     medias: Vec<String>,
     data_urls: HashMap<String, String>,
     tool_calls: Option<MessageContentToolCalls>,
+    query_path: Vec<u64>,
+    subquery_data: Option<(
+        Vec<Message>,
+        Option<f64>,
+        Option<f64>,
+        Option<Vec<crate::function::FunctionDeclaration>>,
+    )>,
     role: Role,
     rag_name: Option<String>,
     with_session: bool,
@@ -47,6 +54,8 @@ impl Input {
             medias: Default::default(),
             data_urls: Default::default(),
             tool_calls: None,
+            query_path: vec![],
+            subquery_data: None,
             role,
             rag_name: None,
             with_session,
@@ -114,6 +123,8 @@ impl Input {
             medias,
             data_urls,
             tool_calls: Default::default(),
+            query_path: vec![],
+            subquery_data: None,
             role,
             rag_name: None,
             with_session,
@@ -167,6 +178,25 @@ impl Input {
         self.config.read().stream && !self.role().model().no_stream()
     }
 
+    pub fn query_path(&self) -> Vec<u64> {
+        self.query_path.clone()
+    }
+
+    pub fn set_query_path(&mut self, path: Vec<u64>) {
+        self.query_path = path;
+    }
+
+    pub fn set_subquery(
+        &mut self,
+        messages: Vec<Message>,
+        temperature: Option<f64>,
+        top_p: Option<f64>,
+        functions: Option<Vec<crate::function::FunctionDeclaration>>,
+    ) {
+        self.subquery_data = Some((messages, temperature, top_p, functions));
+        self.with_session = false;
+    }
+
     pub fn continue_output(&self) -> Option<&str> {
         self.continue_output.as_deref()
     }
@@ -235,6 +265,17 @@ impl Input {
         model: &Model,
         stream: bool,
     ) -> Result<ChatCompletionsData> {
+        if let Some((subquery_messages, subquery_temperature, subquery_top_p, subquery_functions)) =
+            &self.subquery_data
+        {
+            return Ok(crate::client::ChatCompletionsData {
+                messages: subquery_messages.clone(),
+                temperature: *subquery_temperature,
+                top_p: *subquery_top_p,
+                functions: subquery_functions.clone(),
+                stream,
+            });
+        }
         let mut messages = self.build_messages()?;
         patch_messages(&mut messages, model);
         model.guard_max_input_tokens(&messages)?;