Skip to content

Commit 742365f

Browse files
kennvene-awskennvene-aws
committed
[fix] Use Agent instead of Agent Name
* Computes unique identifier based on path and agent name.
1 parent f4ca5d8 commit 742365f

File tree

5 files changed

+65
-36
lines changed

5 files changed

+65
-36
lines changed

‎crates/chat-cli/src/cli/chat/cli/knowledge.rs‎

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -102,9 +102,9 @@ impl KnowledgeSubcommand {
102102
}
103103
}
104104

105-
/// Get the current agent name from the session
106-
fn get_agent_name(session: &ChatSession) -> Option<&str> {
107-
session.conversation.agents.get_active().map(|a| a.name.as_str())
105+
/// Get the current agent from the session
106+
fn get_agent(session: &ChatSession) -> Option<&crate::cli::Agent> {
107+
session.conversation.agents.get_active()
108108
}
109109

110110
async fn execute_operation(&self, os: &Os, session: &mut ChatSession) -> OperationResult {
@@ -132,7 +132,7 @@ impl KnowledgeSubcommand {
132132
}
133133

134134
async fn handle_show(os: &Os, session: &mut ChatSession) -> Result<(), std::io::Error> {
135-
let agent_name = Self::get_agent_name(session).map(|s| s.to_string());
135+
let agent_name = Self::get_agent(session).map(|a| a.name.clone());
136136

137137
// Show agent-specific knowledge
138138
if let Some(agent) = agent_name {
@@ -144,7 +144,7 @@ impl KnowledgeSubcommand {
144144
style::SetAttribute(crossterm::style::Attribute::Reset),
145145
)?;
146146

147-
match KnowledgeStore::get_async_instance(os, Some(agent.as_str())).await {
147+
match KnowledgeStore::get_async_instance(os, Self::get_agent(session)).await {
148148
Ok(store) => {
149149
let store = store.lock().await;
150150
let contexts = store.get_all().await.unwrap_or_default();
@@ -244,9 +244,9 @@ impl KnowledgeSubcommand {
244244
) -> OperationResult {
245245
match Self::validate_and_sanitize_path(os, path) {
246246
Ok(sanitized_path) => {
247-
let agent_name = Self::get_agent_name(session);
247+
let agent = Self::get_agent(session);
248248

249-
let async_knowledge_store = match KnowledgeStore::get_async_instance(os, agent_name).await {
249+
let async_knowledge_store = match KnowledgeStore::get_async_instance(os, agent).await {
250250
Ok(store) => store,
251251
Err(e) => return OperationResult::Error(format!("Error accessing knowledge base: {}", e)),
252252
};
@@ -294,9 +294,9 @@ impl KnowledgeSubcommand {
294294
/// Handle remove operation
295295
async fn handle_remove(os: &Os, session: &ChatSession, path: &str) -> OperationResult {
296296
let sanitized_path = sanitize_path_tool_arg(os, path);
297-
let agent_name = Self::get_agent_name(session);
297+
let agent = Self::get_agent(session);
298298

299-
let async_knowledge_store = match KnowledgeStore::get_async_instance(os, agent_name).await {
299+
let async_knowledge_store = match KnowledgeStore::get_async_instance(os, agent).await {
300300
Ok(store) => store,
301301
Err(e) => return OperationResult::Error(format!("Error accessing knowledge base: {}", e)),
302302
};
@@ -324,8 +324,8 @@ impl KnowledgeSubcommand {
324324
async fn handle_update(os: &Os, session: &ChatSession, path: &str) -> OperationResult {
325325
match Self::validate_and_sanitize_path(os, path) {
326326
Ok(sanitized_path) => {
327-
let agent_name = Self::get_agent_name(session);
328-
let async_knowledge_store = match KnowledgeStore::get_async_instance(os, agent_name).await {
327+
let agent = Self::get_agent(session);
328+
let async_knowledge_store = match KnowledgeStore::get_async_instance(os, agent).await {
329329
Ok(store) => store,
330330
Err(e) => {
331331
return OperationResult::Error(format!("Error accessing knowledge base directory: {}", e));
@@ -363,8 +363,8 @@ impl KnowledgeSubcommand {
363363
return OperationResult::Info("Clear operation cancelled".to_string());
364364
}
365365

366-
let agent_name = Self::get_agent_name(session);
367-
let async_knowledge_store = match KnowledgeStore::get_async_instance(os, agent_name).await {
366+
let agent = Self::get_agent(session);
367+
let async_knowledge_store = match KnowledgeStore::get_async_instance(os, agent).await {
368368
Ok(store) => store,
369369
Err(e) => return OperationResult::Error(format!("Error accessing knowledge base directory: {}", e)),
370370
};
@@ -398,8 +398,8 @@ impl KnowledgeSubcommand {
398398

399399
/// Handle status operation
400400
async fn handle_status(os: &Os, session: &ChatSession) -> OperationResult {
401-
let agent_name = Self::get_agent_name(session);
402-
let async_knowledge_store = match KnowledgeStore::get_async_instance(os, agent_name).await {
401+
let agent = Self::get_agent(session);
402+
let async_knowledge_store = match KnowledgeStore::get_async_instance(os, agent).await {
403403
Ok(store) => store,
404404
Err(e) => return OperationResult::Error(format!("Error accessing knowledge base directory: {}", e)),
405405
};
@@ -510,8 +510,8 @@ impl KnowledgeSubcommand {
510510

511511
/// Handle cancel operation
512512
async fn handle_cancel(os: &Os, session: &ChatSession, operation_id: Option<&str>) -> OperationResult {
513-
let agent_name = Self::get_agent_name(session);
514-
let async_knowledge_store = match KnowledgeStore::get_async_instance(os, agent_name).await {
513+
let agent = Self::get_agent(session);
514+
let async_knowledge_store = match KnowledgeStore::get_async_instance(os, agent).await {
515515
Ok(store) => store,
516516
Err(e) => return OperationResult::Error(format!("Error accessing knowledge base directory: {}", e)),
517517
};

‎crates/chat-cli/src/cli/chat/tools/knowledge.rs‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -315,10 +315,9 @@ impl Knowledge {
315315
&self,
316316
os: &Os,
317317
_updates: &mut impl Write,
318-
agent: Option<&crate::cli::agent::Agent>,
318+
agent: Option<&crate::cli::Agent>,
319319
) -> Result<InvokeOutput> {
320-
let agent_name = agent.map(|a| a.name.as_str());
321-
let async_knowledge_store = KnowledgeStore::get_async_instance(os, agent_name)
320+
let async_knowledge_store = KnowledgeStore::get_async_instance(os, agent)
322321
.await
323322
.map_err(|e| eyre::eyre!("Failed to access knowledge base: {}", e))?;
324323
let mut store = async_knowledge_store.lock().await;

‎crates/chat-cli/src/cli/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ use std::process::ExitCode;
1717

1818
use agent::AgentArgs;
1919
use anstream::println;
20+
pub use agent::{Agent, DEFAULT_AGENT_NAME};
2021
pub use chat::ConversationState;
2122
use clap::{
2223
ArgAction,

‎crates/chat-cli/src/util/directories.rs‎

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use std::path::{
55

66
use thiserror::Error;
77

8+
use crate::cli::DEFAULT_AGENT_NAME;
89
use crate::os::Os;
910

1011
#[derive(Debug, Error)]
@@ -191,8 +192,33 @@ pub fn knowledge_bases_dir(os: &Os) -> Result<PathBuf> {
191192
}
192193

193194
/// The directory for agent-specific knowledge base storage
194-
pub fn agent_knowledge_dir(os: &Os, agent_name: &str) -> Result<PathBuf> {
195-
Ok(knowledge_bases_dir(os)?.join(agent_name))
195+
pub fn agent_knowledge_dir(os: &Os, agent: Option<&crate::cli::Agent>) -> Result<PathBuf> {
196+
let unique_id = if let Some(agent) = agent {
197+
generate_agent_unique_id(agent)
198+
} else {
199+
// Default agent case
200+
DEFAULT_AGENT_NAME.to_string()
201+
};
202+
Ok(knowledge_bases_dir(os)?.join(unique_id))
203+
}
204+
205+
/// Generate a unique identifier for an agent based on its path and name
206+
fn generate_agent_unique_id(agent: &crate::cli::Agent) -> String {
207+
use std::collections::hash_map::DefaultHasher;
208+
use std::hash::{Hash, Hasher};
209+
210+
if let Some(path) = &agent.path {
211+
// Create a hash from the agent's path for uniqueness
212+
let mut hasher = DefaultHasher::new();
213+
path.hash(&mut hasher);
214+
let path_hash = hasher.finish();
215+
216+
// Combine hash with agent name for readability
217+
format!("{}_{:x}", agent.name, path_hash)
218+
} else {
219+
// For agents without a path (like default), use just the name
220+
agent.name.clone()
221+
}
196222
}
197223

198224
/// The path to the fig settings file

‎crates/chat-cli/src/util/knowledge_store.rs‎

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,10 @@ use semantic_search_client::types::{
1515
use tokio::sync::Mutex;
1616
use uuid::Uuid;
1717

18+
use crate::cli::DEFAULT_AGENT_NAME;
1819
use crate::os::Os;
1920
use crate::util::directories;
2021

21-
const DEFAULT_AGENT_NAME: &str = "q_cli_default";
22-
2322
/// Configuration for adding knowledge contexts
2423
#[derive(Default)]
2524
pub struct AddOptions {
@@ -111,31 +110,30 @@ pub struct KnowledgeStore {
111110
}
112111

113112
impl KnowledgeStore {
114-
/// Get singleton instance with optional agent name
113+
/// Get singleton instance with optional agent
115114
pub async fn get_async_instance(
116115
os: &Os,
117-
agent_name: Option<&str>,
116+
agent: Option<&crate::cli::Agent>,
118117
) -> Result<Arc<Mutex<Self>>, directories::DirectoryError> {
119118
static ASYNC_INSTANCE: Lazy<tokio::sync::OnceCell<Arc<Mutex<KnowledgeStore>>>> =
120119
Lazy::new(tokio::sync::OnceCell::new);
121120

122121
if cfg!(test) {
123122
// For tests, create a new instance each time
124-
let store = Self::new_with_os_settings(os, agent_name)
123+
let store = Self::new_with_os_settings(os, agent)
125124
.await
126125
.map_err(|_e| directories::DirectoryError::Io(std::io::Error::other("Failed to create store")))?;
127126
Ok(Arc::new(Mutex::new(store)))
128127
} else {
129128
Ok(ASYNC_INSTANCE
130129
.get_or_init(|| async {
131130
// Check for migration before initializing the client
132-
let agent_name = agent_name.unwrap_or(DEFAULT_AGENT_NAME);
133-
let agent_dir = crate::util::directories::agent_knowledge_dir(os, agent_name)
131+
let agent_dir = crate::util::directories::agent_knowledge_dir(os, agent)
134132
.expect("Failed to get agent directory");
135133

136-
Self::migrate_legacy_knowledge_base(&agent_dir, agent_name).await;
134+
Self::migrate_legacy_knowledge_base(&agent_dir).await;
137135

138-
let store = Self::new_with_os_settings(os, Some(agent_name))
136+
let store = Self::new_with_os_settings(os, agent)
139137
.await
140138
.expect("Failed to create knowledge store");
141139
Arc::new(Mutex::new(store))
@@ -146,8 +144,14 @@ impl KnowledgeStore {
146144
}
147145

148146
/// Migrate legacy knowledge base from old location if needed
149-
async fn migrate_legacy_knowledge_base(agent_dir: &PathBuf, current_agent_name: &str) -> bool {
147+
async fn migrate_legacy_knowledge_base(agent_dir: &PathBuf) -> bool {
150148
let mut migrated = false;
149+
150+
// Extract agent identifier from the directory path (last component)
151+
let current_agent_id = agent_dir
152+
.file_name()
153+
.and_then(|name| name.to_str())
154+
.unwrap_or(DEFAULT_AGENT_NAME);
151155

152156
// Migrate from legacy ~/.semantic_search
153157
let old_flat_dir = dirs::home_dir()
@@ -177,7 +181,7 @@ impl KnowledgeStore {
177181
.filter(|entry| {
178182
let name = entry.file_name();
179183
let name_str = name.to_string_lossy();
180-
name_str != current_agent_name
184+
name_str != current_agent_id
181185
&& name_str != DEFAULT_AGENT_NAME
182186
&& !name_str.starts_with('.')
183187
})
@@ -245,9 +249,8 @@ impl KnowledgeStore {
245249
}
246250

247251
/// Create instance with database settings from OS
248-
async fn new_with_os_settings(os: &crate::os::Os, agent_name: Option<&str>) -> Result<Self> {
249-
let agent_name = agent_name.unwrap_or(DEFAULT_AGENT_NAME);
250-
let agent_dir = crate::util::directories::agent_knowledge_dir(os, agent_name)?;
252+
async fn new_with_os_settings(os: &crate::os::Os, agent: Option<&crate::cli::Agent>) -> Result<Self> {
253+
let agent_dir = crate::util::directories::agent_knowledge_dir(os, agent)?;
251254
let agent_config = Self::create_config_from_db_settings(os, agent_dir.clone());
252255
let agent_client = AsyncSemanticSearchClient::with_config(&agent_dir, agent_config)
253256
.await

0 commit comments

Comments
 (0)