axonerai 0.1.0

A type-safe, blazing fast agentic AI framework 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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

use crate::executor::{ToolExecutor, ToolResult};
use crate::provider::{Message, Provider, StopReason};
use crate::tool::ToolRegistry;
use anyhow::Result;
use crate::file_session_manager::FileSessionManager;
use crate::session::Session;

pub struct Agent {
    provider: Box<dyn Provider>,
    registry: ToolRegistry,
    max_iterations: usize,
    system_prompt: Option<String>,
    file_session_manager: Option<FileSessionManager>
}

impl Agent {
    pub fn new(provider: Box<dyn Provider>, registry: ToolRegistry, system_prompt: Option<String>, file_session_manager: Option<FileSessionManager>) -> Self {
        Self {
            provider,
            registry,
            max_iterations: 10, // Prevent infinite loops
            system_prompt,
            file_session_manager
        }
    }

    /// Run the agent with a user prompt
    pub async fn run(&self, user_prompt: &str) -> Result<String> {

        let mut session = if let Some(ref sm) = self.file_session_manager {
            if sm.exists() { sm.load()? }
            else { Session::new(sm.get_session().to_string()) }
        }
        else {
            Session::new("stateless".to_string())
        };



        println!();

        session.add_message(Message {
            role: "user".to_string(),
            content: user_prompt.to_string(),
        });

        let executor = ToolExecutor::new(&self.registry);
        let tools = self.registry.get_all_for_llm();

        for _iteration in 1..=self.max_iterations {

            let response = self
                .provider
                .complete(session.get_messages().clone(), Some(tools.clone()), None, self.system_prompt.clone())
                .await?;

            match response.stop_reason {
                StopReason::EndTurn => {
                    if let Some(text) = response.text {
                        session.add_message(Message {
                            role: "assistant".to_string(),
                            content: text.clone(),
                        });

                        if let Some(ref sm) = self.file_session_manager {
                            sm.save(&session)?;
                        }

                        println!("Response from Agent:");
                        return Ok(text);
                    } else {
                        return Ok("(No response from agent)".to_string());
                    }
                }

                StopReason::ToolUse => {
                    // LLM wants to use tools
                    println!("Entered here");
                    if let Some(text) = &response.text {
                        println!("💭 Agent thinking: {}", text);
                    }

                    if response.tool_calls.is_empty() {
                        return Ok("Agent wanted to use tools but didn't specify any".to_string());
                    }

                    // Execute the tools
                    let tool_results = executor.execute_all(&response.tool_calls).await?;

                    // Add assistant's tool use to messages
                    session.add_message(Message {
                        role: "assistant".to_string(),
                        content: format_tool_use(&response.tool_calls),
                    });

                    // Add tool results to messages
                    session.add_message(Message {
                        role: "user".to_string(),
                        content: format_tool_results(&tool_results),
                    });

                    println!();
                    // Continue the loop
                }

                StopReason::MaxTokens => {
                    return Ok("Agent hit max tokens limit".to_string());
                }

                _ => {
                    return Ok(format!("Agent stopped with reason: {:?}", response.stop_reason));
                }
            }
        }

        if let Some(ref sm) = self.file_session_manager {
            sm.save(&session)?;
        }
        Ok(format!(
            "Agent reached max iterations ({})",
            self.max_iterations
        ))
    }
}

fn format_tool_use(tool_calls: &[crate::provider::ToolCall]) -> String {
    tool_calls
        .iter()
        .map(|call| {
            format!(
                "Using tool '{}' with input: {}",
                call.name,
                serde_json::to_string(&call.input).unwrap_or_default()
            )
        })
        .collect::<Vec<_>>()
        .join("\n")
}

/// Format tool results for feeding back to the LLM
fn format_tool_results(results: &[ToolResult]) -> String {
    results
        .iter()
        .map(|result| format!("Tool '{}' returned: {}", result.tool_name, result.result))
        .collect::<Vec<_>>()
        .join("\n")
}