Skip to content

Commit

Permalink
feat: switch to HashSet for system resource support check
Browse files Browse the repository at this point in the history
  • Loading branch information
kalvinnchau committed Jan 18, 2025
1 parent fcec628 commit 4b1425f
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 5 deletions.
8 changes: 4 additions & 4 deletions crates/goose/src/agents/capabilities.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use chrono::{DateTime, TimeZone, Utc};
use futures::stream::{FuturesUnordered, StreamExt};
use mcp_client::McpService;
use rust_decimal_macros::dec;
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::sync::LazyLock;
use std::time::Duration;
Expand All @@ -26,7 +26,7 @@ static DEFAULT_TIMESTAMP: LazyLock<DateTime<Utc>> =
pub struct Capabilities {
clients: HashMap<String, Arc<Mutex<Box<dyn McpClientTrait>>>>,
instructions: HashMap<String, String>,
resource_capable_systems: Vec<String>,
resource_capable_systems: HashSet<String>,
provider: Box<dyn Provider>,
provider_usage: Mutex<Vec<ProviderUsage>>,
}
Expand Down Expand Up @@ -84,7 +84,7 @@ impl Capabilities {
Self {
clients: HashMap::new(),
instructions: HashMap::new(),
resource_capable_systems: Vec::new(),
resource_capable_systems: HashSet::new(),
provider,
provider_usage: Mutex::new(Vec::new()),
}
Expand Down Expand Up @@ -153,7 +153,7 @@ impl Capabilities {
// if the server is capable if resources we track it
if init_result.capabilities.resources.is_some() {
self.resource_capable_systems
.push(init_result.server_info.name.clone());
.insert(init_result.server_info.name.clone());
}

// Store the client
Expand Down
2 changes: 1 addition & 1 deletion crates/goose/src/agents/reference.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use futures::stream::BoxStream;
use tokio::sync::Mutex;
use tracing::{debug, instrument};

use super::{capabilities, Agent};
use super::{Agent};
use crate::agents::capabilities::Capabilities;
use crate::agents::system::{SystemConfig, SystemResult};
use crate::message::{Message, ToolRequest};
Expand Down

0 comments on commit 4b1425f

Please sign in to comment.