diff options
| author | aichat-agent <agent@aichat.local> | 2026-10-01 14:58:06 +0000 |
|---|---|---|
| committer | aichat-agent <agent@aichat.local> | 2026-10-01 15:54:19 +0000 |
| commit | 83a2cd1ac24ca0feda34ae877aaf31f41b58b86c (patch) | |
| tree | cbc83678fa05fe6a82908c448a4e20eae673a9c7 /src/utils | |
| parent | 82976d349ad97ac9aae0655ad631dace5e2a6385 (diff) | |
| download | aichat-83a2cd1ac24ca0feda34ae877aaf31f41b58b86c.tar.gz | |
feat: add built-in subquery tool with hierarchical query ids
* Add config options subquery and subquery_max_depth.
* Register an internal subquery function that spawns a fresh LLM query
with the current model configuration but without the full context.
The model selects context numerically (message_indexes) or by text
excerpt (excerpts).
* Subqueries can use tools, including further nested subqueries up to
subquery_max_depth.
* Queries and subqueries get hierarchical ids ([0], [0/1], [0/1/2])
that prefix tool-call prints and spinner messages.
Diffstat (limited to 'src/utils')
| -rw-r--r-- | src/utils/abort_signal.rs | 68 |
1 files changed, 67 insertions, 1 deletions
diff --git a/src/utils/abort_signal.rs b/src/utils/abort_signal.rs index a13e5bb..47534ae 100644 --- a/src/utils/abort_signal.rs +++ b/src/utils/abort_signal.rs @@ -2,7 +2,7 @@ use anyhow::Result; use crossterm::event::{self, Event, KeyCode, KeyModifiers}; use std::{ sync::{ - atomic::{AtomicBool, Ordering}, + atomic::{AtomicBool, AtomicU64, Ordering}, Arc, }, time::Duration, @@ -10,6 +10,72 @@ use std::{ pub type AbortSignal = Arc<AbortSignalInner>; +#[derive(Debug, Default)] +pub struct QueryIdGen { + counter: AtomicU64, + parents: parking_lot::RwLock<std::collections::HashMap<String, u64>>, +} + +impl QueryIdGen { + pub fn root_id(&self) -> u64 { + self.counter.fetch_add(1, Ordering::SeqCst) + } + + pub fn child_id(&self, parent_path: &[u64]) -> u64 { + if parent_path.is_empty() { + return self.root_id(); + } + let key = parent_path + .iter() + .map(|v| v.to_string()) + .collect::<Vec<_>>() + .join("/"); + let mut parents = self.parents.write(); + let entry = parents.entry(key).or_insert(0); + let id = *entry; + *entry += 1; + id + } +} + +pub fn format_query_path(path: &[u64]) -> String { + if path.is_empty() { + return String::new(); + } + format!( + "[{}]", + path.iter() + .map(|v| v.to_string()) + .collect::<Vec<_>>() + .join("/") + ) +} + +tokio::task_local! { + static QUERY_SCOPE: String; + static QUERY_ID_GEN: QueryIdGen; +} + +pub fn current_query_scope() -> Option<String> { + QUERY_SCOPE + .try_with(|v| if v.is_empty() { None } else { Some(v.clone()) }) + .ok() + .flatten() +} + +pub async fn with_query_scope<F: std::future::Future>(scope: String, future: F) -> F::Output { + QUERY_SCOPE.scope(scope, future).await +} + +pub fn next_query_id(parent_path: &[u64]) -> u64 { + if parent_path.is_empty() { + return QUERY_ID_GEN.try_with(|g| g.root_id()).unwrap_or(0); + } + QUERY_ID_GEN + .try_with(|g| g.child_id(parent_path)) + .unwrap_or(0) +} + pub struct AbortSignalInner { ctrlc: AtomicBool, ctrld: AtomicBool, |
