lm_studio_api_extended/chat/
context.rs

1// use crate::prelude::*;
2use super::{ Role, Message };
3
4// Chat context
5#[derive(Debug, Clone)]
6pub struct Context {
7    pub messages: Vec<Message>,
8    pub context_tokens: usize,
9    pub context_limit: usize
10}
11
12impl Default for Context {
13    fn default() -> Self {
14        Context { 
15            messages: vec![
16                Message { 
17                    role: Role::System, 
18                    content: "You are a helpful, knowledgeable, and friendly assistant.".to_string()
19                }
20            ], 
21            context_tokens: 0, 
22            context_limit: 4090 
23        }
24    }
25}
26
27impl Context {
28    // Creates a new chat context
29    pub fn new<S>(context: S, context_limit: usize) -> Self
30    where
31        S: Into<String>
32    {
33        Self {
34            messages: vec![Message::new(Role::System, context.into())],
35            context_tokens: 0,
36            context_limit
37        }
38    }
39
40    // Adds context from RAG to system prompt
41    pub fn edit<S>(&mut self, modification: S)
42    where 
43        S: Into<String>
44    {
45        let sysprompt_extra = format!("\n\nContext: [\n\t{}]", modification.into());
46        self.messages[0].content.push_str(&sysprompt_extra);
47    }
48
49    // Add a message to context
50    pub fn add<M>(&mut self, message: M)
51    where
52        M: Into<Message>
53    {
54        let message = message.into();
55
56        // add message to context:
57        self.context_tokens += message.content.chars().count();
58        self.messages.push(message);
59
60        // remove old extra messages:
61        while self.messages.len() > 2 && self.context_tokens > self.context_limit {
62            self.context_tokens -= self.messages[1].content.chars().count();
63            self.messages.remove(1);
64        }
65    }
66
67    // Get all context messages
68    pub fn get(&self) -> Vec<Message> {
69        self.messages.clone()
70    }
71
72    // Clear context messages
73    pub fn clear(&mut self) {
74        self.messages.truncate(1);
75    }
76}
lm_studio_api_extended/chat/context.rs

lm_studio_api_extended/chat/
context.rs
