class SessionMemory:
    def __init__(self):
        self.sessions = {}  # {session_key: [exchanges]}

    def ensure_session(self, session_key):
        if session_key not in self.sessions:
            self.sessions[session_key] = []

    def add_exchange(self, session_key, rep_input, persona_response, persona_question, objection, traits_snapshot):
        self.ensure_session(session_key)
        self.sessions[session_key].append({
            "rep_input": rep_input,
            "persona_response": persona_response,
            "persona_question": persona_question,
            "objection": objection,
            "traits_snapshot": traits_snapshot
        })

    def get_history(self, session_key):
        self.ensure_session(session_key)
        return self.sessions.get(session_key, [])

    def reset(self):
        self.sessions = {}