Pythonmemory.py
Copy
1
from openai import OpenAI
2
import tiktoken
3
from typing import Literal
4
 
5
client = OpenAI()
6
enc = tiktoken.encoding_for_model("gpt-4o")
7
 
8
class SlidingWindowMemory:
9
    def __init__(self, max_tokens: int = 8000, system_reserve: int = 600):
10
        self.messages: list[dict] = []
11
        self.max_tokens = max_tokens
12
        self.system_reserve = system_reserve
13
 
14
    @property
15
    def budget(self) -> int:
16
        return self.max_tokens - self.system_reserve
17
 
18
    def _count(self) -> int:
19
        return sum(4 + len(enc.encode(m["content"])) for m in self.messages)
20
 
21
    def add(self, role: Literal["user", "assistant"], content: str):
22
        self.messages.append({"role": role, "content": content})
23
        self._trim()
24
 
25
    def _trim(self):
26
        while self._count() > self.budget and len(self.messages) > 2:
27
            self.messages.pop(0)
28
 
29
    def context(self) -> list[dict]:
30
        return self.messages.copy()
31
 
32
    def stats(self) -> dict:
33
        return {"messages": len(self.messages), "tokens": self._count(), "budget": self.budget}
34
 
35
memory = SlidingWindowMemory(max_tokens=4000)
36
 
37
def chat(user_input: str) -> str:
38
    memory.add("user", user_input)
39
    response = client.chat.completions.create(
40
        model="gpt-4o",
41
        messages=[{"role": "system", "content": "You are a helpful assistant."}] + memory.context()
42
    )
43
    answer = response.choices[0].message.content
44
    memory.add("assistant", answer)
45
    print(f"Memory stats: {memory.stats()}")
46
    return answer
47
 
48
if __name__ == "__main__":
49
    print(chat("My name is Alice and I'm learning about LLMs."))
50
    print(chat("What is my name?"))
51
    print(chat("What are we discussing?"))
52