lm_studio_api/chat/
context.rs

1// use crate::prelude::*;
2use super::{ Role, Message };
3
4// Chat context
5#[derive(Debug, Clone)]
6pub struct Context {
7    messages: Vec<Message>,
8    context_tokens: usize,
9    context_limit: usize
10}
11
12impl Context {
13    // Creates a new chat context
14    pub fn new<S>(context: S, context_limit: usize) -> Self
15    where
16        S: Into<String>
17    {
18        Self {
19            messages: vec![Message::new(Role::System, context.into())],
20            context_tokens: 0,
21            context_limit
22        }
23    }
24
25    // Add a message to context
26    pub fn add<M>(&mut self, message: M)
27    where
28        M: Into<Message>
29    {
30        let message = message.into();
31
32        // add message to context:
33        self.context_tokens += message.content.chars().count();
34        self.messages.push(message);
35
36        // remove old extra messages:
37        while self.messages.len() > 2 && self.context_tokens > self.context_limit {
38            self.context_tokens -= self.messages[1].content.chars().count();
39            self.messages.remove(1);
40        }
41    }
42
43    // Get all context messages
44    pub fn get(&self) -> Vec<Message> {
45        self.messages.clone()
46    }
47
48    // Clear context messages
49    pub fn clear(&mut self) {
50        self.messages.truncate(1);
51    }
52}