Building a Session Manager

Welcome back! In the previous lesson, we explored the importance of a robust system prompt and how it guides the behavior of our personal tutor. Now, we will delve into the next step of our journey: building the Session Manager. This lesson will focus on the model layer of the MVC (Model-View-Controller) architecture, which is crucial for organizing and managing data in a structured way.

The SessionManager struct will be the core component of our model layer, responsible for managing tutoring session data effectively. By the end of this lesson, you will understand how to create, manage, and retrieve session data using the SessionManager struct.

The SessionManager struct is designed to handle the storage and management of tutoring session data. It serves as the backbone of our tutor's data management system. We'll start by setting up the struct and then gradually add methods to handle session creation, message addition, and conversation retrieval.

Let's begin by defining the SessionManager struct and its constructor. The constructor initializes an empty HashMap, sessions, which will store all session data.

use std::collections::HashMap;

pub struct SessionData {
    pub system_prompt: String,
    pub messages: Vec<ChatCompletionRequestMessage>,
}

pub struct SessionManager {
    sessions: HashMap<String, HashMap<String, SessionData>>,
}

impl SessionManager {
    pub fn new() -> Self {
        SessionManager { sessions: HashMap::new() }
    }
}

In this setup, sessions is a nested HashMap where the first key is the student_id, and the second key is the session_id. This structure allows us to efficiently manage multiple tutoring sessions for different students.

Next, we'll add the create_session method. This method is responsible for creating a new session entry for a student. It takes three parameters: student_id, session_id, and system_prompt.

pub fn create_session(
    &mut self,
    student_id: impl Into<String>,
    session_id: impl Into<String>,
    system_prompt: impl Into<String>,
) {
    let sid = student_id.into();
    let sess = session_id.into();
    let prompt = system_prompt.into();
    let data = SessionData {
        system_prompt: prompt,
        messages: Vec::new(),
    };
    self.sessions
        .entry(sid)
        .or_default()
        .insert(sess, data);
}

The create_session method checks if the student_id exists in sessions. If not, it creates a new entry. Then, it initializes the session with the provided system_prompt and an empty vector for messages.

If a session with the same session_id already exists for a student, this method will overwrite the existing session without warning. In production, you may want to check for existing sessions first and return an error or log a warning to prevent accidental data loss.

To access a specific tutoring session, we need a method to retrieve a session based on the student_id and session_id. In Rust, we use the Option type and HashMap access patterns for this.

pub(crate) fn get_session_mut(
    &mut self,
    student_id: &str,
    session_id: &str,
) -> Option<&mut SessionData> {
    self.sessions
        .get_mut(student_id)
        .and_then(|m| m.get_mut(session_id))
}

The get_session_mut method returns a mutable reference to the session data if it exists, or None otherwise.

Now, let's add the add_message method. This method allows us to append messages to a session. It requires the student_id, session_id, role, and content of the message.

pub fn add_message(
    &mut self,
    student_id: &str,
    session_id: &str,
    role: Role,
    content: &str,
) -> Result<()> {
    let session = self
        .get_session_mut(student_id, session_id)
        .context("Session not found")?;

    let msg = match role {
        Role::System => ChatCompletionRequestSystemMessageArgs::default()
            .content(content)
            .build()?
            .into(),
        Role::User => ChatCompletionRequestUserMessageArgs::default()
            .content(content)
            .build()?
            .into(),
        Role::Assistant => ChatCompletionRequestAssistantMessageArgs::default()
            .content(content)
            .build()?
            .into(),
        _ => bail!("Unsupported role: {:?}", role),
    };

    session.messages.push(msg);
    Ok(())
}

The add_message method first retrieves the session using get_session_mut. If the session exists, it appends the message to the session's message vector.

Finally, we'll implement the get_conversation method. This method returns the entire conversation, including the system prompt and all messages.

pub fn get_conversation(
    &self,
    student_id: &str,
    session_id: &str,
) -> Result<Vec<ChatCompletionRequestMessage>> {
    let session = self
        .sessions
        .get(student_id)
        .and_then(|m| m.get(session_id))
        .context("Session not found")?;

    let mut convo = Vec::with_capacity(1 + session.messages.len());

    let system_msg = ChatCompletionRequestSystemMessageArgs::default()
        .content(session.system_prompt.as_str())
        .build()?
        .into();

    convo.push(system_msg);
    convo.extend(session.messages.clone());

    Ok(convo)
}

The get_conversation method retrieves the session and constructs a vector starting with the system prompt, followed by the session messages. If the session does not exist, it returns an error.

Note that this method clones the session’s messages. While this is acceptable for small conversations, cloning large vectors can impact performance. For high-throughput systems or sessions with large histories, consider using reference-based access patterns (such as returning a &[ChatCompletionRequestMessage]) to avoid cloning.

To create a new tutoring session, we use the create_session method. This method takes three parameters: student_id, session_id, and system_prompt. It initializes a new session entry in the sessions HashMap. If the student_id does not already exist in the map, a new entry is created. The session entry includes the system prompt and an empty vector for messages.

Here's an example of how to create a new tutoring session:

let mut manager = SessionManager::new();
let system_prompt = load_system_prompt("data/system_prompt.txt");
manager.create_session("student_test", "session_test", system_prompt);

In this example, we initialize a SessionManager instance and create a new tutoring session for a student with the ID student_test. The session is identified by session_test and is initialized with a system prompt. This setup allows us to manage tutoring session data efficiently.

Once a session is created, we can add messages to it using the add_message method. This method requires the student_id, session_id, role, and content of the message. The role indicates whether the message is from the student or the tutor assistant. The message is then appended to the session's message vector.

Here's how you can add messages to a tutoring session:

manager.add_message("student_test", "session_test", Role::User, "Hello, I need help with geometry.")?;
manager.add_message("student_test", "session_test", Role::Assistant, "Hi there! What specific geometry topic are you struggling with?")?;

In this example, we add a student message, Hello, I need help with geometry., and an assistant response, Hi there! What specific geometry topic are you struggling with?, to the session. The add_message method ensures that messages are stored in the correct session entry.

To retrieve the entire conversation, including the system prompt, we use the get_conversation method. This method returns a vector of messages, starting with the system prompt followed by the session messages.

// Retrieve the full conversation
let conversation = manager.get_conversation(student_id, session_id)?;

// Display the session details
println!("Session ID: {}\n", session_id);
for message in &conversation {
    println!("Role: {:?}", message.role);
    println!("Message:\n{}\n", message.content.as_deref().unwrap_or(""));
}

The get_conversation method compiles the session history, allowing us to access the full context of the conversation. This is particularly useful when we need to provide the full conversation history to the model to ensure contextual responses.

In this lesson, we explored the SessionManager struct and its role in managing tutoring session data within the model layer of our application. We learned how to create and manage sessions, add messages, and retrieve conversation histories. The SessionManager is a crucial component for organizing session data, ensuring that our personal tutor can handle multiple tutoring sessions efficiently.

As you move on to the practice exercises, take the opportunity to experiment with modifying and extending the SessionManager functionality. This hands-on practice will reinforce the concepts covered in this lesson and prepare you for the next steps in our course. Keep up the great work, and I look forward to seeing your progress!