research-papers-rag / memory.py
akzz
added error handling and logs
2baf97f
Raw
History Blame Contribute Delete
4.39 kB
"""
Manages in-memory conversation history for a session
Resets when app restarts
Structure:
memory = {
"summary": str, # compresses older convo turns
"recent": list[dict] # has {"role":"...", "content":"..."} of latest few turns
}
When recent turns exceed certain number (RECENT_TURNS), the older turns are summarised and folded into 'summary', keeping the prompt size bounded without lossing track of earlier context.
What this file does:
1. initiate new empty memory
2. summarize some older turns and updates the summary
3. add new turns of conversation in memory, and if the turn's number exceed the threshold, get summary
4. function to properly format the history or memory to pass to the prompt for LLM
"""
import os
from dotenv import load_dotenv
from google import genai
from config import LLM_MODEL
from logger import get_logger
load_dotenv()
logger = get_logger("memory")
llm = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
RECENT_TURNS = 6 # number of full (user+assistant) turn-pairs to keep
COMPRESS_BY = 3 # when overflow, fold this many oldest pairs into the summary
# 1. initiate new empty memory
def new_memory() -> dict:
"""Return a blank memory object"""
return {"summary":"", "recent": []}
# 2. summarize some older turns and updates the summary
def summarise(existing_summary: str, turns_to_compress: list) -> str:
"""
Ask the LLM to fold `turns_to_compress` into `existing_summary`.
Returns the updated summary string.
"""
try:
turns_text = "\n".join(
f"{'User' if t["role"] == 'user' else 'Assistant'}:{t['content']}"
for t in turns_to_compress
)
prior_block = (
f"Existing summary:\n{existing_summary}\n\n" if existing_summary else ""
)
prompt = (
f"{prior_block}"
f"New conversation turns to add to the summary: \n{turns_text}\n\n"
"Write a concise but complete summary of the full conversation so far."
"Preserve key facts, questions asked, and answers given."
"Return ONLY the summary, nothing else."
)
response = llm.models.generate_content(
model=LLM_MODEL,
contents=prompt
)
new_summary = response.text.strip()
logger.info("summarise: compressed %d turns into summary", len(turns_to_compress))
return new_summary
except Exception as e:
logger.warning("summarise: LLM call failed, keeping existing summary: %s", e)
return existing_summary # don't lose old summary if compression fails
# 3. add new turns of conversation in memory, and if the turn's number exceed the threshold, get summary
def add_turn(memory: dict, role:str, content:str) -> dict:
"""
Append one message to memory.
If recent turns exceed the limit, compress the oldest ones into the summary.
Always returns the updated memory dict.
"""
memory["recent"].append({"role":role, "content":content})
max_messages = RECENT_TURNS * 2 # pairs to individual messages
if len(memory["recent"]) > max_messages:
compress_count = COMPRESS_BY * 2 # messages to compress
to_compress = memory["recent"][:compress_count] # take the starting compress_count messages
memory["recent"] = memory["recent"][compress_count:] # reassign the recent turns, starting from the compress_count nth turn to latest turn
logger.info("add_turn: compressing %d messages into summary", compress_count)
memory["summary"] = summarise(memory["summary"], to_compress)
return memory
# 4. function to properly format the history or memory to pass to the prompt for LLM
def format_history_for_prompt(memory: dict) -> str:
"""
Renders the full memory (summary + recent turns) as a plain-text block
ready to inject into the LLM prompt.
Returns an empty string if memory is empty.
"""
parts = []
if memory.get("summary"):
parts.append(f"Summary of earlier conversation:\n{memory['summary']}")
if memory.get("recent"):
recent_lines = []
for turn in memory["recent"]:
label = "User" if turn["role"] == "user" else "Assistant"
recent_lines.append(f"{label}:{turn['content']}")
parts.append("Recent conversation:\n" + "\n".join(recent_lines))
return "\n\n".join(parts)