1from openai import OpenAI
2from typing import Literal
3import tiktoken
4
5client = OpenAI()
6enc = tiktoken.encoding_for_model("gpt-4o")
7
8MAX_TOKENS = 8000
9SYSTEM_TOKENS = 500
10
11class SlidingWindowMemory:
12 def __init__(self, max_tokens: int = MAX_TOKENS):
13 self.messages: list[dict] = []
14 self.max_tokens = max_tokens
15
16 def count_tokens(self, messages: list[dict]) -> int:
17 return sum(len(enc.encode(m["content"])) for m in messages)
18
19 def add(self, role: Literal["user", "assistant"], content: str):
20 self.messages.append({"role": role, "content": content})
21 self._trim()
22
23 def _trim(self):
24 available = self.max_tokens - SYSTEM_TOKENS
25 while self.count_tokens(self.messages) > available and len(self.messages) > 2:
26 self.messages.pop(0) # Remove oldest message
27
28 def compress(self):
29 if len(self.messages) < 6:
30 return
31 old_messages = self.messages[:-4]
32 summary = client.chat.completions.create(
33 model="gpt-4o-mini",
34 messages=[{"role": "user", "content": f"Summarise this conversation in 2 sentences:
35
36" + "
37".join(f"{m['role']}: {m['content']}" for m in old_messages)}]
38 ).choices[0].message.content
39 self.messages = [{"role": "user", "content": f"[Summary of earlier conversation: {summary}]"}] + self.messages[-4:]
40
41 def get_context(self) -> list[dict]:
42 return self.messages
43
44memory = SlidingWindowMemory()
45def chat(user_input: str) -> str:
46 memory.add("user", user_input)
47 response = client.chat.completions.create(
48 model="gpt-4o",
49 messages=[{"role": "system", "content": "You are a helpful assistant."}] + memory.get_context()
50 )
51 answer = response.choices[0].message.content
52 memory.add("assistant", answer)
53 return answer
