use super::expertise::Expertise;
use super::{Agent, AgentError, Capability, Payload, ToExpertise};
use crate::prompt::{PromptPart, ToPrompt};
use async_trait::async_trait;
pub struct ExpertiseAgent<T: Agent> {
inner_agent: T,
expertise: Expertise,
}
impl<T: Agent> ExpertiseAgent<T> {
pub fn new(inner_agent: T, expertise: Expertise) -> Self {
Self {
inner_agent,
expertise,
}
}
pub fn inner(&self) -> &T {
&self.inner_agent
}
pub fn expertise_ref(&self) -> &Expertise {
&self.expertise
}
}
#[async_trait]
impl<T> Agent for ExpertiseAgent<T>
where
T: Agent + Send + Sync,
T::Output: Send,
{
type Output = T::Output;
type Expertise = Expertise;
fn expertise(&self) -> &Expertise {
&self.expertise
}
async fn execute(&self, intent: Payload) -> Result<Self::Output, AgentError> {
let render_context = intent.render_context().cloned().unwrap_or_default();
let expertise_prompt = self.expertise.to_prompt_with_context(&render_context);
let enriched_payload = intent.prepend_system(expertise_prompt);
self.inner_agent.execute(enriched_payload).await
}
fn name(&self) -> String {
format!("ExpertiseAgent<{}>", self.inner_agent.name())
}
async fn is_available(&self) -> Result<(), AgentError> {
self.inner_agent.is_available().await
}
}
impl<T: Agent> ToPrompt for ExpertiseAgent<T> {
fn to_prompt_parts(&self) -> Vec<PromptPart> {
self.expertise.to_prompt_parts()
}
fn to_prompt(&self) -> String {
self.expertise.to_prompt()
}
}
impl<T: Agent> ToExpertise for ExpertiseAgent<T> {
fn description(&self) -> &str {
self.expertise.description()
}
fn capabilities(&self) -> Vec<Capability> {
self.expertise.capabilities()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::agent::expertise::{KnowledgeFragment, WeightedFragment};
use crate::context::{ContextProfile, TaskHealth};
struct MockAgent {
response: String,
}
#[async_trait]
impl Agent for MockAgent {
type Output = String;
type Expertise = &'static str;
fn expertise(&self) -> &&'static str {
&"Mock agent"
}
async fn execute(&self, _intent: Payload) -> Result<Self::Output, AgentError> {
Ok(self.response.clone())
}
fn name(&self) -> String {
"MockAgent".to_string()
}
async fn is_available(&self) -> Result<(), AgentError> {
Ok(())
}
}
#[tokio::test]
async fn test_expertise_agent_basic() {
let mock = MockAgent {
response: "Mock response".to_string(),
};
let expertise = Expertise::new("test", "1.0").with_fragment(WeightedFragment::new(
KnowledgeFragment::Text("Test content".to_string()),
));
let agent = ExpertiseAgent::new(mock, expertise);
let payload = Payload::text("Question");
let result = agent.execute(payload).await.unwrap();
assert_eq!(result, "Mock response");
}
#[tokio::test]
async fn test_expertise_agent_with_context() {
let mock = MockAgent {
response: "Mock response".to_string(),
};
let expertise = Expertise::new("test", "1.0")
.with_fragment(
WeightedFragment::new(KnowledgeFragment::Text("Always visible".to_string()))
.with_context(ContextProfile::Always),
)
.with_fragment(
WeightedFragment::new(KnowledgeFragment::Text("AtRisk only".to_string()))
.with_context(ContextProfile::Conditional {
task_types: vec![],
user_states: vec![],
task_health: Some(TaskHealth::AtRisk),
}),
);
let agent = ExpertiseAgent::new(mock, expertise);
let payload = Payload::text("Question").with_task_health(TaskHealth::AtRisk);
let result = agent.execute(payload).await.unwrap();
assert_eq!(result, "Mock response");
}
#[tokio::test]
async fn test_expertise_agent_name() {
let mock = MockAgent {
response: "Mock response".to_string(),
};
let expertise = Expertise::new("test", "1.0");
let agent = ExpertiseAgent::new(mock, expertise);
assert_eq!(agent.name(), "ExpertiseAgent<MockAgent>");
}
#[tokio::test]
async fn test_expertise_agent_is_available() {
let mock = MockAgent {
response: "Mock response".to_string(),
};
let expertise = Expertise::new("test", "1.0");
let agent = ExpertiseAgent::new(mock, expertise);
assert!(agent.is_available().await.is_ok());
}
#[tokio::test]
async fn test_expertise_agent_inner_access() {
let mock = MockAgent {
response: "Mock response".to_string(),
};
let expertise = Expertise::new("test", "1.0");
let agent = ExpertiseAgent::new(mock, expertise);
assert_eq!(agent.inner().name(), "MockAgent");
}
#[tokio::test]
async fn test_expertise_agent_expertise_ref() {
let mock = MockAgent {
response: "Mock response".to_string(),
};
let expertise = Expertise::new("test-expertise", "1.0");
let agent = ExpertiseAgent::new(mock, expertise);
assert_eq!(agent.expertise_ref().id, "test-expertise");
}
#[test]
fn test_expertise_agent_to_expertise() {
let mock = MockAgent {
response: "Mock response".to_string(),
};
let expertise =
Expertise::new("test-desc", "1.0").with_description("Test description for routing");
let agent = ExpertiseAgent::new(mock, expertise);
assert_eq!(
ToExpertise::description(&agent),
"Test description for routing"
);
assert_eq!(ToExpertise::capabilities(&agent).len(), 0);
}
#[test]
fn test_expertise_agent_to_prompt() {
use crate::agent::expertise::{KnowledgeFragment, WeightedFragment};
let mock = MockAgent {
response: "Mock response".to_string(),
};
let expertise = Expertise::new("test", "1.0").with_fragment(WeightedFragment::new(
KnowledgeFragment::Text("Test content".to_string()),
));
let agent = ExpertiseAgent::new(mock, expertise);
let prompt = agent.to_prompt();
assert!(prompt.contains("Test content"));
}
}