hehe_agent/
lib.rs

1pub mod error;
2pub mod config;
3pub mod event;
4pub mod session;
5pub mod response;
6pub mod executor;
7pub mod agent;
8
9pub use error::{AgentError, Result};
10pub use config::AgentConfig;
11pub use event::AgentEvent;
12pub use session::{Session, SessionStats};
13pub use response::{AgentResponse, ToolCallRecord};
14pub use agent::{Agent, AgentBuilder};
15
16pub mod prelude {
17    pub use crate::error::{AgentError, Result};
18    pub use crate::config::AgentConfig;
19    pub use crate::event::AgentEvent;
20    pub use crate::session::{Session, SessionStats};
21    pub use crate::response::{AgentResponse, ToolCallRecord};
22    pub use crate::agent::{Agent, AgentBuilder};
23}
24
25#[cfg(test)]
26mod tests {
27    use super::*;
28    use async_trait::async_trait;
29    use hehe_core::capability::Capabilities;
30    use hehe_core::stream::StreamChunk;
31    use hehe_core::Message;
32    use hehe_llm::{BoxStream, CompletionRequest, CompletionResponse, LlmError, LlmProvider, ModelInfo};
33    use std::sync::Arc;
34
35    struct MockLlm;
36
37    #[async_trait]
38    impl LlmProvider for MockLlm {
39        fn name(&self) -> &str {
40            "mock"
41        }
42
43        fn capabilities(&self) -> &Capabilities {
44            static CAPS: std::sync::OnceLock<Capabilities> = std::sync::OnceLock::new();
45            CAPS.get_or_init(Capabilities::text_basic)
46        }
47
48        async fn complete(&self, _request: CompletionRequest) -> std::result::Result<CompletionResponse, LlmError> {
49            Ok(CompletionResponse::new("id", "mock", Message::assistant("Test response")))
50        }
51
52        async fn complete_stream(
53            &self,
54            _request: CompletionRequest,
55        ) -> std::result::Result<BoxStream<StreamChunk>, LlmError> {
56            use futures::stream;
57            Ok(Box::pin(stream::empty()))
58        }
59
60        async fn list_models(&self) -> std::result::Result<Vec<ModelInfo>, LlmError> {
61            Ok(vec![])
62        }
63
64        fn default_model(&self) -> &str {
65            "mock"
66        }
67    }
68
69    #[tokio::test]
70    async fn test_full_agent_flow() {
71        let agent = Agent::builder()
72            .system_prompt("You are a helpful assistant.")
73            .model("mock")
74            .llm(Arc::new(MockLlm))
75            .build()
76            .unwrap();
77
78        let session = agent.create_session();
79
80        let response = agent.chat(&session, "Hello!").await.unwrap();
81        assert_eq!(response, "Test response");
82
83        let response2 = agent.chat(&session, "How are you?").await.unwrap();
84        assert_eq!(response2, "Test response");
85
86        assert_eq!(session.message_count(), 4);
87
88        let stats = session.stats();
89        assert_eq!(stats.message_count, 4);
90        assert_eq!(stats.iteration_count, 2);
91    }
92
93    #[tokio::test]
94    async fn test_agent_with_config() {
95        let config = AgentConfig::new("mock", "System prompt")
96            .with_name("test-agent")
97            .with_max_iterations(5)
98            .with_temperature(0.3);
99
100        let agent = Agent::builder()
101            .config(config)
102            .llm(Arc::new(MockLlm))
103            .build()
104            .unwrap();
105
106        assert_eq!(agent.config().name, "test-agent");
107        assert_eq!(agent.config().max_iterations, 5);
108        assert_eq!(agent.config().temperature, 0.3);
109    }
110}
hehe_agent/lib.rs

hehe_agent/
lib.rs
