agent-base 0.1.0

A lightweight Agent Runtime Kernel for building AI agents in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

use std::collections::HashSet;

use serde::{Deserialize, Serialize};

use crate::types::{ChatMessage, ImageAttachment, Message, MessageRole, ToolCallMessage};

use crate::types::SessionId;

#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct AgentSession {
    id: Option<SessionId>,
    messages: Vec<Message>,
    chat_messages: Vec<ChatMessage>,
    always_allowed_actions: HashSet<String>,
}

impl AgentSession {
    pub fn new(id: SessionId) -> Self {
        Self {
            id: Some(id),
            messages: Vec::new(),
            chat_messages: Vec::new(),
            always_allowed_actions: HashSet::new(),
        }
    }

    pub fn id(&self) -> Option<SessionId> {
        self.id.clone()
    }

    pub fn messages(&self) -> &[Message] {
        &self.messages
    }

    pub fn chat_messages(&self) -> &[ChatMessage] {
        &self.chat_messages
    }

    pub fn is_action_allowed(&self, action_key: &str) -> bool {
        self.always_allowed_actions.contains(action_key)
    }

    pub fn allow_action(&mut self, action_key: impl Into<String>) {
        self.always_allowed_actions.insert(action_key.into());
    }

    pub fn push_message(&mut self, role: MessageRole, content: impl Into<String>) {
        let content = content.into();
        self.messages.push(Message {
            role: role.clone(),
            content: content.clone(),
        });
        let chat_msg = match role {
            MessageRole::System => ChatMessage::system(content),
            MessageRole::User => ChatMessage::user(content),
            MessageRole::Assistant => ChatMessage::assistant(content),
            MessageRole::Tool => ChatMessage::tool(String::new(), content),
        };
        self.chat_messages.push(chat_msg);
    }

    pub fn push_user_message_with_images(
        &mut self,
        content: impl Into<String>,
        images: Vec<ImageAttachment>,
    ) {
        let content = content.into();
        self.messages.push(Message {
            role: MessageRole::User,
            content: content.clone(),
        });
        self.chat_messages
            .push(ChatMessage::user_with_images(content, images));
    }

    pub fn push_assistant_tool_call(
        &mut self,
        tool_call_id: &str,
        tool_name: &str,
        arguments_json: &str,
    ) {
        self.chat_messages.push(ChatMessage::assistant_tool_call(
            tool_call_id,
            tool_name,
            arguments_json,
        ));
    }

    pub fn push_assistant_tool_calls(&mut self, tool_calls: &[(String, String, String)]) {
        let calls: Vec<ToolCallMessage> = tool_calls
            .iter()
            .map(|(id, name, args)| ToolCallMessage {
                id: id.clone(),
                name: name.clone(),
                arguments: args.clone(),
            })
            .collect();
        self.chat_messages.push(ChatMessage::Assistant {
            content: None,
            reasoning_content: None,
            tool_calls: Some(calls),
        });
    }

    pub fn push_tool_result(&mut self, tool_call_id: &str, content: impl Into<String>) {
        let content = content.into();
        self.messages.push(Message {
            role: MessageRole::Tool,
            content: content.clone(),
        });
        self.chat_messages.push(ChatMessage::tool(
            tool_call_id,
            content,
        ));
    }
}