Struct version of tool & knowledge (#212)

Author: jhlee525

v2/src/agent/agent.rs

@@ -5,9 +5,9 @@ use futures::{Stream, StreamExt, lock::Mutex};
 
 use crate::{
     agent::SystemMessageRenderer,
-    knowledge::Knowledge,
+    knowledge::{Knowledge, KnowledgeBehavior as _},
     model::{InferenceConfig, LangModel, LangModelInference as _},
-    tool::Tool,
+    tool::{Tool, ToolBehavior as _},
     utils::log,
     value::{Delta, FinishReason, Message, MessageDelta, Part, PartDelta, Role},
     warn,
@@ -16,9 +16,9 @@ use crate::{
 #[derive(Clone)]
 pub struct Agent {
     lm: LangModel,
-    tools: Vec<Arc<dyn Tool>>,
+    tools: Vec<Tool>,
     messages: Arc<Mutex<Vec<Message>>>,
-    knowledge: Option<Arc<dyn Knowledge>>,
+    knowledge: Option<Knowledge>,
     system_message_renderer: Arc<SystemMessageRenderer>,
 }
 
@@ -34,7 +34,7 @@ pub struct AgentResponse {
 }
 
 impl Agent {
-    pub fn new(lm: LangModel, tools: impl IntoIterator<Item = Arc<dyn Tool>>) -> Self {
+    pub fn new(lm: LangModel, tools: impl IntoIterator<Item = Tool>) -> Self {
         Self {
             lm,
             tools: tools.into_iter().collect(),
@@ -55,11 +55,11 @@ impl Agent {
         self.lm.clone()
     }
 
-    pub fn get_tools(&self) -> Vec<Arc<dyn Tool>> {
+    pub fn get_tools(&self) -> Vec<Tool> {
         self.tools.clone()
     }
 
-    pub fn knowledge(&self) -> Option<Arc<dyn Knowledge>> {
+    pub fn knowledge(&self) -> Option<Knowledge> {
         self.knowledge.clone()
     }
 
@@ -68,7 +68,7 @@ impl Agent {
         Ok(())
     }
 
-    pub async fn add_tools(&mut self, tools: Vec<Arc<dyn Tool>>) -> anyhow::Result<()> {
+    pub async fn add_tools(&mut self, tools: Vec<Tool>) -> anyhow::Result<()> {
         for tool in tools.iter() {
             let tool_name = tool.get_description().name;
 
@@ -92,7 +92,7 @@ impl Agent {
         Ok(())
     }
 
-    pub async fn add_tool(&mut self, tool: Arc<dyn Tool>) -> anyhow::Result<()> {
+    pub async fn add_tool(&mut self, tool: Tool) -> anyhow::Result<()> {
         self.add_tools(vec![tool]).await
     }
 
@@ -109,17 +109,8 @@ impl Agent {
         self.remove_tools(vec![tool_name]).await
     }
 
-    pub async fn remove_mcp_tools(&mut self, client_name: String) -> anyhow::Result<()> {
-        self.tools.retain(|t| {
-            let tool_name = t.get_description().name;
-            // Remove the MCP tool if its description name is prefixed with the provided client name.
-            !tool_name.starts_with(format!("{}--", client_name).as_str())
-        });
-        Ok(())
-    }
-
-    pub fn set_knowledge(&mut self, knowledge: impl Knowledge + 'static) {
-        self.knowledge = Some(Arc::new(knowledge));
+    pub fn set_knowledge(&mut self, knowledge: Knowledge) {
+        self.knowledge = Some(knowledge);
     }
 
     pub fn remove_knowledge(&mut self) {
@@ -140,10 +131,10 @@ impl Agent {
             let system_message_content = "You are helpful assistant.".to_string();
             let knowledge_results = if let Some(knowledge) = &self.knowledge {
                 let query = contents.iter().filter(|&c| matches!(c, Part::Text{..})).map(|c| c.as_text().unwrap()).collect::<Vec<_>>().join("\n");
-                let retrieved = match knowledge.retrieve(query.clone()).await {
+                let retrieved = match knowledge.retrieve(query.clone(), 1).await {
                     Ok(retrieved) => retrieved,
                     Err(e) => {
-                        warn!("Failed to retrieve from knowledge {}: {}", knowledge.name(), e.to_string());
+                        warn!("Failed to retrieve from knowledge: {}", e.to_string());
                         vec![]
                     }
                 };
@@ -319,7 +310,7 @@ mod tests {
     //                 Part::Text("104".to_owned())
     //             }
     //         }),
-    //     )) as Arc<dyn Tool>];
+    //     )) as Tool];
     //     let mut agent = Agent::new(model, tools);
 
     //     let mut agg = MessageAggregator::new();

v2/src/knowledge/base.rs

@@ -4,9 +4,12 @@ use ailoy_macros::{maybe_send_sync, multi_platform_async_trait};
 use serde::{Deserialize, Serialize};
 
 use crate::{
+    knowledge::{CustomKnowledge, VectorStoreKnowledge},
+    model::EmbeddingModel,
     to_value,
-    tool::Tool,
+    tool::ToolBehavior,
     value::{ToolDesc, Value},
+    vector_store::VectorStore,
 };
 
 type Metadata = serde_json::Map<String, serde_json::Value>;
@@ -19,15 +22,17 @@ pub struct KnowledgeRetrieveResult {
 
 #[maybe_send_sync]
 #[multi_platform_async_trait]
-pub trait Knowledge: std::fmt::Debug {
-    fn name(&self) -> String;
-
-    async fn retrieve(&self, query: String) -> anyhow::Result<Vec<KnowledgeRetrieveResult>>;
+pub trait KnowledgeBehavior: std::fmt::Debug {
+    async fn retrieve(
+        &self,
+        query: String,
+        top_k: u32,
+    ) -> anyhow::Result<Vec<KnowledgeRetrieveResult>>;
 }
 
 #[derive(Clone)]
 pub struct KnowledgeTool {
-    inner: Arc<dyn Knowledge>,
+    inner: Arc<dyn KnowledgeBehavior>,
     desc: ToolDesc,
 }
 
@@ -36,15 +41,14 @@ impl std::fmt::Debug for KnowledgeTool {
         f.debug_struct("KnowledgeTool")
             .field("desc", &self.desc)
             .field("inner", &self.inner)
-            .field("stringify", &"(Function)")
             .finish()
     }
 }
 
 impl KnowledgeTool {
-    pub fn from(knowledge: impl Knowledge + 'static) -> Self {
+    pub fn from(knowledge: impl KnowledgeBehavior + 'static) -> Self {
         let default_desc = ToolDesc {
-            name: format!("retrieve-{}", knowledge.name()),
+            name: "retrieve-from-knowledge".into(),
             description: Some("Retrieve the relevant context from knowledge base.".into()),
             parameters: to_value!({
                 "type": "object",
@@ -71,7 +75,7 @@ impl KnowledgeTool {
 }
 
 #[multi_platform_async_trait]
-impl Tool for KnowledgeTool {
+impl ToolBehavior for KnowledgeTool {
     fn get_description(&self) -> ToolDesc {
         self.desc.clone()
     }
@@ -94,7 +98,7 @@ impl Tool for KnowledgeTool {
             }
         };
 
-        let results = match self.inner.retrieve(query.into()).await {
+        let results = match self.inner.retrieve(query.into(), 1).await {
             Ok(results) => results,
             Err(e) => {
                 return Ok(e.to_string().into());
@@ -106,58 +110,44 @@ impl Tool for KnowledgeTool {
     }
 }
 
-#[cfg(test)]
-mod tests {
-    use ailoy_macros::multi_platform_test;
-    use futures::stream::StreamExt;
-
-    use super::*;
-    use crate::{agent::Agent, model::LangModel, value::Part};
+#[derive(Debug, Clone)]
+pub enum KnowledgeInner {
+    VectorStore(VectorStoreKnowledge),
+    Custom(CustomKnowledge),
+}
 
-    #[derive(Debug)]
-    struct CustomKnowledge {}
+#[derive(Debug, Clone)]
+pub struct Knowledge {
+    inner: KnowledgeInner,
+}
 
-    #[multi_platform_async_trait]
-    impl Knowledge for CustomKnowledge {
-        fn name(&self) -> String {
-            "about-ailoy".into()
+impl Knowledge {
+    pub fn new_vector_store(
+        store: impl VectorStore + 'static,
+        embedding_model: EmbeddingModel,
+    ) -> Self {
+        Self {
+            inner: KnowledgeInner::VectorStore(VectorStoreKnowledge::new(store, embedding_model)),
         }
+    }
 
-        async fn retrieve(&self, _query: String) -> anyhow::Result<Vec<KnowledgeRetrieveResult>> {
-            let documents = vec![
-                KnowledgeRetrieveResult {
-                    document: "Ailoy is an awesome AI agent framework.".into(),
-                    metadata: None,
-                },
-                KnowledgeRetrieveResult {
-                    document: "Ailoy supports Python, Javascript and Rust.".into(),
-                    metadata: None,
-                },
-                KnowledgeRetrieveResult {
-                    document: "Ailoy enables running LLMs in local environment easily.".into(),
-                    metadata: None,
-                },
-            ];
-            Ok(documents)
+    pub fn new_custom(knowledge: CustomKnowledge) -> Self {
+        Self {
+            inner: KnowledgeInner::Custom(knowledge),
         }
     }
+}
 
-    #[multi_platform_test]
-    async fn test_custom_knowledge_with_agent() -> anyhow::Result<()> {
-        let knowledge = CustomKnowledge {};
-        let model = LangModel::try_new_local("Qwen/Qwen3-0.6B").await.unwrap();
-        let mut agent = Agent::new(model, vec![]);
-
-        agent.set_knowledge(knowledge);
-
-        let mut strm = Box::pin(agent.run(vec![Part::Text {
-            text: "What is Ailoy?".into(),
-        }]));
-        while let Some(out) = strm.next().await {
-            let out = out.unwrap();
-            println!("{:?}", out);
+#[multi_platform_async_trait]
+impl KnowledgeBehavior for Knowledge {
+    async fn retrieve(
+        &self,
+        query: String,
+        top_k: u32,
+    ) -> anyhow::Result<Vec<KnowledgeRetrieveResult>> {
+        match &self.inner {
+            KnowledgeInner::VectorStore(knowledge) => knowledge.retrieve(query, top_k).await,
+            KnowledgeInner::Custom(knowledge) => knowledge.retrieve(query, top_k).await,
         }
-
-        Ok(())
     }
 }

v2/src/knowledge/custom_knowledge.rs

@@ -0,0 +1,96 @@
+use std::{fmt::Debug, sync::Arc};
+
+use ailoy_macros::multi_platform_async_trait;
+use futures::future::BoxFuture;
+
+use crate::{
+    knowledge::{KnowledgeBehavior, KnowledgeRetrieveResult},
+    utils::{MaybeSend, MaybeSync},
+};
+
+#[derive(Clone)]
+pub struct CustomKnowledge {
+    f: Arc<
+        dyn Fn(String, u32) -> BoxFuture<'static, anyhow::Result<Vec<KnowledgeRetrieveResult>>>
+            + MaybeSend
+            + MaybeSync,
+    >,
+}
+
+impl Debug for CustomKnowledge {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("CustomKnowledge")
+            .field("f", &"function")
+            .finish()
+    }
+}
+
+impl CustomKnowledge {
+    pub fn new(
+        f: Arc<
+            dyn Fn(String, u32) -> BoxFuture<'static, anyhow::Result<Vec<KnowledgeRetrieveResult>>>
+                + MaybeSend
+                + MaybeSync,
+        >,
+    ) -> Self {
+        Self { f }
+    }
+}
+
+#[multi_platform_async_trait]
+impl KnowledgeBehavior for CustomKnowledge {
+    async fn retrieve(
+        &self,
+        query: String,
+        top_k: u32,
+    ) -> anyhow::Result<Vec<KnowledgeRetrieveResult>> {
+        (self.f)(query, top_k).await
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use ailoy_macros::multi_platform_test;
+    use futures::{FutureExt, stream::StreamExt};
+
+    use super::*;
+    use crate::{agent::Agent, knowledge::Knowledge, model::LangModel, value::Part};
+
+    #[multi_platform_test]
+    async fn test_custom_knowledge_with_agent() -> anyhow::Result<()> {
+        let knowledge = Knowledge::new_custom(CustomKnowledge::new(Arc::new(|_, _| {
+            async {
+                let documents = vec![
+                    KnowledgeRetrieveResult {
+                        document: "Ailoy is an awesome AI agent framework.".into(),
+                        metadata: None,
+                    },
+                    KnowledgeRetrieveResult {
+                        document: "Ailoy supports Python, Javascript and Rust.".into(),
+                        metadata: None,
+                    },
+                    KnowledgeRetrieveResult {
+                        document: "Ailoy enables running LLMs in local environment easily.".into(),
+                        metadata: None,
+                    },
+                ];
+                Ok(documents)
+            }
+            .boxed()
+        })));
+        let model = LangModel::try_new_local("Qwen/Qwen3-0.6B").await.unwrap();
+        let mut agent = Agent::new(model, vec![]);
+
+        agent.set_knowledge(knowledge);
+
+        let mut strm = Box::pin(agent.run(vec![Part::Text {
+            text: "What is Ailoy?".into(),
+        }]));
+        while let Some(out) = strm.next().await {
+            let out = out.unwrap();
+            println!("{:?}", out);
+        }
+
+        Ok(())
+    }
+}

v2/src/knowledge/mod.rs

@@ -1,5 +1,7 @@
 mod base;
+mod custom_knowledge;
 mod vector_store_knowledge;
 
 pub use base::*;
+pub use custom_knowledge::*;
 pub use vector_store_knowledge::*;