Short-Term Memory Basics

python -m venv .venv && source .venv/bin/activate
pip install openai
export OPENAI_API_KEY="sk-..."
# Optional gateway:
# export OPENAI_BASE_URL="https://openrouter.ai/api/v1"

from openai import OpenAI
 
client = OpenAI()  # uses OPENAI_API_KEY; set base_url for gateways
model = "gpt-4o-mini"  # verify at build
 
r1 = client.chat.completions.create(
    model=model,
    messages=[{"role": "user", "content": "My name is Ada. Reply in one short sentence."}],
)
print(r1.choices[0].message.content)
 
r2 = client.chat.completions.create(
    model=model,
    messages=[{"role": "user", "content": "What is my name?"}],
)
print(r2.choices[0].message.content)  # model has no prior turn

messages = [
    {"role": "system", "content": "You are a concise assistant. Remember user facts stated in this chat."},
]
 
def say(user_text: str) -> str:
    messages.append({"role": "user", "content": user_text})
    resp = client.chat.completions.create(model=model, messages=messages)
    text = resp.choices[0].message.content or ""
    messages.append({"role": "assistant", "content": text})
    return text
 
print(say("My name is Ada."))
print(say("What is my name?"))

def show_history():
    for i, m in enumerate(messages):
        content = m.get("content") or ""
        print(f"{i}: {m['role']:10} {len(content):4} chars  {content[:60]!r}")
 
show_history()

def trim_messages(msgs: list, max_messages: int = 12) -> list:
    system = [m for m in msgs if m["role"] == "system"]
    rest = [m for m in msgs if m["role"] != "system"]
    if len(rest) > max_messages:
        rest = rest[-max_messages:]
    return system + rest
 
# After each turn:
# messages[:] = trim_messages(messages, max_messages=12)

from typing import Any
 
session: dict[str, Any] = {
    "id": "sess_demo_1",
    "messages": [
        {"role": "system", "content": "Help the user complete a small checklist."},
    ],
    "scratchpad": {
        "goal": None,
        "steps_done": [],
    },
}
 
def update_goal(goal: str) -> None:
    session["scratchpad"]["goal"] = goal
    session["scratchpad"]["steps_done"] = []
 
update_goal("Plan a 3-item grocery list")
session["messages"].append({
    "role": "user",
    "content": f"Goal: {session['scratchpad']['goal']}. Suggest three items.",
})
resp = client.chat.completions.create(model=model, messages=session["messages"])
session["messages"].append({"role": "assistant", "content": resp.choices[0].message.content})
session["scratchpad"]["steps_done"].append("suggested_items")
print(session["scratchpad"])

from collections import defaultdict
 
# In-process demo only; use Redis or DB in production
STORE: dict[tuple[str, str, str], dict] = defaultdict(lambda: {
    "messages": [{"role": "system", "content": "You are helpful and brief."}],
    "scratchpad": {},
})
 
def get_session(tenant_id: str, user_id: str, session_id: str) -> dict:
    return STORE[(tenant_id, user_id, session_id)]
 
a = get_session("acme", "u1", "s1")
b = get_session("acme", "u2", "s1")
a["messages"].append({"role": "user", "content": "secret-from-u1"})
assert b["messages"][-1]["content"] != "secret-from-u1" if len(b["messages"]) > 1 else True

import json
from pathlib import Path
 
path = Path("/tmp/agent_session_demo.json")
path.write_text(json.dumps(session), encoding="utf-8")
 
restored = json.loads(path.read_text(encoding="utf-8"))
print(restored["id"], len(restored["messages"]))

def messages_for_model(sess: dict) -> list:
    pad = json.dumps(sess["scratchpad"], ensure_ascii=False)
    extra = {
        "role": "system",
        "content": f"Current scratchpad (authoritative JSON):\n{pad}",
    }
    # Keep base system + scratchpad projection + dialogue
    base = [m for m in sess["messages"] if m["role"] == "system"][:1]
    dialogue = [m for m in sess["messages"] if m["role"] != "system"]
    return base + [extra] + dialogue
 
# resp = client.chat.completions.create(model=model, messages=messages_for_model(session))