State persistence across sessions in Agentic AI - ML Experiment: Train & Evaluate

Experiment - State persistence across sessions

Problem:You have an AI agent that learns from user interactions during a session. However, when the session ends, the agent loses all learned information and starts fresh next time.

Current Metrics:Session 1: Agent accuracy 85%, Session 2: Agent accuracy resets to 50% (random guess).

Issue:The agent does not save its learned state between sessions, causing loss of progress and poor performance in new sessions.

Your Task

Implement state persistence so the agent retains learned knowledge across sessions, improving accuracy in subsequent sessions to at least 80%.

You must keep the agent's learning algorithm unchanged.

You can only add mechanisms to save and load the agent's state.

Use simple file-based storage for persistence.

Hint 1

Hint 2

Hint 3

Solution

Agentic AI

import pickle

class SimpleAgent:
    def __init__(self):
        self.knowledge = {}

    def learn(self, data):
        for item in data:
            self.knowledge[item] = self.knowledge.get(item, 0) + 1

    def predict(self, item):
        return self.knowledge.get(item, 0) > 0

    def save_state(self, filename):
        with open(filename, 'wb') as f:
            pickle.dump(self.knowledge, f)

    def load_state(self, filename):
        try:
            with open(filename, 'rb') as f:
                self.knowledge = pickle.load(f)
        except FileNotFoundError:
            self.knowledge = {}


# Simulate session 1
agent = SimpleAgent()
agent.load_state('agent_state.pkl')
training_data_session1 = ['apple', 'banana', 'apple', 'orange']
agent.learn(training_data_session1)
agent.save_state('agent_state.pkl')

# Test prediction after session 1
print('Session 1 prediction for apple:', agent.predict('apple'))  # Expected: True

# Simulate session 2
agent2 = SimpleAgent()
agent2.load_state('agent_state.pkl')

# Test prediction at start of session 2
print('Session 2 prediction for apple:', agent2.predict('apple'))  # Expected: True

# Learn new data in session 2
training_data_session2 = ['banana', 'banana', 'grape']
agent2.learn(training_data_session2)
agent2.save_state('agent_state.pkl')

# Test prediction after session 2 learning
print('Session 2 prediction for grape:', agent2.predict('grape'))  # Expected: True

Added methods save_state and load_state to save and load the agent's knowledge dictionary using pickle.

Loaded saved state at the start of each session to retain learned knowledge.

Saved updated state at the end of each session to persist new learning.

Results Interpretation

Before: Each new session started with no knowledge, causing accuracy to drop to random guess levels (~50%).

After: The agent loads saved knowledge at session start, maintaining learned information and achieving over 80% accuracy in subsequent sessions.

Saving and loading the agent's state allows it to remember past learning, improving performance across sessions and demonstrating the importance of state persistence in AI systems.

Bonus Experiment

Try implementing state persistence using a database instead of file storage to handle larger or multiple agents.

💡 Hint

Use a simple key-value store like SQLite or Redis to save and retrieve the agent's knowledge dictionary.