ambi 0.2.4

A flexible, multi-backend, customizable AI agent framework, entirely based on 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

// src/agent/pipeline/chat_runner/context_handler.rs

use super::StateManager;
use crate::agent::core::{Agent, EvictionHandler};
use crate::agent::tool::ToolDefinition;
use crate::error::Result;
use crate::llm::ChatTemplate;
use crate::types::config::EvictionStrategy;
use crate::types::message::{ContentPart, Message};
use crate::types::request::LLMRequest;

impl<'a> StateManager<'a> {
    pub async fn push_user_message(&self, parts: Vec<ContentPart>, tokens: usize) -> Result<()> {
        self.0
            .write()
            .await
            .chat_history
            .push(Message::User { content: parts }, tokens);
        Ok(())
    }

    pub async fn push_tool_message(
        &self,
        content: String,
        tool_id: Option<String>,
        tokens: usize,
    ) -> Result<()> {
        self.0
            .write()
            .await
            .chat_history
            .push(Message::Tool { content, tool_id }, tokens);
        Ok(())
    }

    pub async fn get_snapshot_len(&self) -> Result<usize> {
        Ok(self.0.read().await.chat_history.len())
    }

    pub async fn truncate(&self, len: usize) -> Result<()> {
        self.0.write().await.chat_history.truncate(len);
        Ok(())
    }

    pub async fn get_llm_request(
        &self,
        system_prompt: &str,
        tpl: &ChatTemplate,
        tools: &[ToolDefinition],
        cached_tool_prompt: &str,
        tool_tags: (String, String),
    ) -> Result<LLMRequest> {
        let lock = self.0.read().await;
        Ok(Agent::get_llm_request(
            &lock,
            system_prompt,
            tpl,
            tools,
            cached_tool_prompt,
            tool_tags,
        ))
    }

    pub async fn get_system_overhead(&self) -> Result<usize> {
        Ok(self
            .0
            .read()
            .await
            .chat_history
            .all()
            .iter()
            .filter(|(m, _)| matches!(**m, Message::System { .. }))
            .map(|(_, t)| *t)
            .sum())
    }

    #[allow(clippy::too_many_arguments)]
    pub async fn append_assistant_message_and_evict(
        &self,
        content: String,
        tool_calls: Vec<(String, serde_json::Value, String)>,
        tokens: usize,
        handler: &Option<EvictionHandler>,
        eviction_strategy: &EvictionStrategy,
        prompt_overhead: usize,
    ) -> Result<usize> {
        let mut lock = self.0.write().await;
        lock.chat_history.push(
            Message::Assistant {
                content,
                tool_calls,
            },
            tokens,
        );

        let evicted_msgs = lock.chat_history.evict_old_messages(
            eviction_strategy.keep_head,
            eviction_strategy.keep_tail,
            eviction_strategy.max_safe_tokens,
            prompt_overhead,
        );

        let count = evicted_msgs.len();
        if count > 0 {
            log::debug!("Context truncation: Evicted {} messages.", count);
            if let Some(h) = handler {
                h(evicted_msgs);
            }
        }
        Ok(count)
    }
}