348 lines
14 KiB
Python
348 lines
14 KiB
Python
import os
|
||
import json
|
||
import time
|
||
import datetime
|
||
import re
|
||
from .core_module import CoreModule, GameTrajectory, Observation
|
||
|
||
class MemoryModule(CoreModule):
|
||
"""
|
||
A lightweight memory module:
|
||
1. stores the most recent N turns in a GameTrajectory deque.
|
||
2. synthesises reflections with an LLM.
|
||
"""
|
||
|
||
def __init__(self,
|
||
model_name: str = "claude-3-7-sonnet-latest",
|
||
cache_dir: str = "cache",
|
||
reflection_system_prompt: str = "",
|
||
reflection_prompt: str = "",
|
||
summary_system_prompt: str = "",
|
||
summary_prompt: str = "",
|
||
max_memory: int = 10,
|
||
token_limit: int = 100000,
|
||
use_reflection: bool = True,
|
||
use_summary: bool = False,
|
||
vllm_url=None,
|
||
modal_url=None):
|
||
|
||
print(f"memory module token limit: {token_limit}")
|
||
|
||
super().__init__(
|
||
module_name="memory_module",
|
||
model_name=model_name,
|
||
system_prompt=reflection_system_prompt,
|
||
prompt=reflection_prompt,
|
||
token_limit=token_limit,
|
||
cache_dir=cache_dir,
|
||
vllm_url=vllm_url,
|
||
modal_url=modal_url
|
||
)
|
||
|
||
self.max_memory = max_memory
|
||
self.use_reflection = use_reflection
|
||
self.use_summary = use_summary
|
||
self.summary_system_prompt = summary_system_prompt
|
||
self.summary_prompt = summary_prompt
|
||
self.current_summary = "" # Store the current summary
|
||
|
||
def _load_trajectory(self) -> None:
|
||
"""Load and return trajectory entries (as already‑stringified lines) from disk."""
|
||
|
||
trajectory = GameTrajectory(max_length=self.max_memory)
|
||
if os.path.exists(self.module_file):
|
||
try:
|
||
with open(self.module_file, "r") as f:
|
||
entries = json.load(f)
|
||
|
||
# keep only the last maxlen lines and push them into the deque
|
||
for e in entries[-self.max_memory:]:
|
||
# expect the entry to have been stored as a ready‑to‑print line
|
||
if isinstance(e, str):
|
||
trajectory.add(e)
|
||
# Check if this entry contains a summary and restore it
|
||
if e.startswith("##TRAJECTORY SUMMARY\n"):
|
||
summary_content = e.replace("##TRAJECTORY SUMMARY\n", "").strip()
|
||
if summary_content:
|
||
self.current_summary = summary_content
|
||
except Exception as exc:
|
||
print(f"[MemoryModule] failed to load trajectory: {exc}")
|
||
else:
|
||
print("trajectory entries do not exist.")
|
||
|
||
return trajectory
|
||
|
||
def _append_to_log(self, line: str) -> None:
|
||
"""
|
||
Persist *just the printable line* per update.
|
||
That keeps the on‑disk structure flat and forward‑compatible.
|
||
"""
|
||
try:
|
||
if os.path.exists(self.module_file):
|
||
with open(self.module_file, "r") as f:
|
||
data = json.load(f)
|
||
else:
|
||
data = []
|
||
|
||
data.append(line)
|
||
with open(self.module_file, "w") as f:
|
||
json.dump(data[-self.max_memory:], f, indent=2)
|
||
except Exception as exc:
|
||
print(f"[MemoryModule] failed to write log: {exc}")
|
||
|
||
def _reflect(self,
|
||
prev_context: str,
|
||
current_state: str) -> str:
|
||
"""
|
||
Ask the LLM to write a reflection given the running context string.
|
||
"""
|
||
formatted_prompt = self.prompt.format(
|
||
prev_context=prev_context or "None",
|
||
current_observation=current_state,
|
||
)
|
||
raw = self.api_manager.text_only_completion(
|
||
model_name=self.model_name,
|
||
system_prompt=self.system_prompt,
|
||
prompt=formatted_prompt,
|
||
thinking=False,
|
||
reasoning_effort=self.reasoning_effort,
|
||
token_limit=self.token_limit,
|
||
)
|
||
# returned API response should be a tuple
|
||
actual_raw_text = raw[0]
|
||
# extract "reflection:" section if present
|
||
m = re.search(
|
||
r'(?:^|\n)(?:#\s*)?reflection:(.+?)(?=\n(?:#\s*)?[a-zA-Z]+:|$)',
|
||
actual_raw_text, # Use the extracted text
|
||
re.DOTALL | re.IGNORECASE,
|
||
)
|
||
return (m.group(1).strip() if m else actual_raw_text.strip()) or "No valid reflection produced."
|
||
|
||
def _summarize(self, game_trajectory: str) -> str:
|
||
"""
|
||
Generate a summary of the game trajectory when it exceeds max_memory length.
|
||
"""
|
||
if not self.summary_prompt or not self.use_summary:
|
||
return ""
|
||
|
||
formatted_prompt = self.summary_prompt.format(
|
||
game_trajectory=game_trajectory,
|
||
previous_summary=self.current_summary or "No previous summary."
|
||
)
|
||
|
||
try:
|
||
print(f"[MemoryModule] Generating summary...")
|
||
|
||
raw = self.api_manager.text_only_completion(
|
||
model_name=self.model_name,
|
||
system_prompt=self.summary_system_prompt,
|
||
prompt=formatted_prompt,
|
||
thinking=False,
|
||
reasoning_effort=self.reasoning_effort,
|
||
token_limit=self.token_limit,
|
||
)
|
||
|
||
# returned API response should be a tuple
|
||
actual_raw_text = raw[0] if raw and len(raw) > 0 else ""
|
||
|
||
# Clean and validate the response
|
||
summary = actual_raw_text.strip() if actual_raw_text else ""
|
||
|
||
# Check if we got a valid summary (not empty and not an error message)
|
||
if summary and len(summary) > 10 and "no valid summary" not in summary.lower():
|
||
print(f"[MemoryModule] Successfully generated summary. Length: {len(summary)} chars")
|
||
return summary
|
||
else:
|
||
print(f"[MemoryModule] Generated invalid summary: '{summary[:100]}...'")
|
||
# Return a basic fallback summary
|
||
fallback_summary = f"FALLBACK SUMMARY: Game trajectory contained {len(game_trajectory)} characters of gameplay data. Previous summary: {self.current_summary[:200] if self.current_summary else 'None'}..."
|
||
return fallback_summary
|
||
|
||
except Exception as e:
|
||
print(f"[MemoryModule] Error generating summary: {e}")
|
||
# Return a basic fallback summary
|
||
fallback_summary = f"FALLBACK SUMMARY: Game trajectory contained {len(game_trajectory)} characters of gameplay data. Previous summary: {self.current_summary[:200] if self.current_summary else 'None'}..."
|
||
print(f"[MemoryModule] Using fallback summary due to error.")
|
||
return fallback_summary
|
||
|
||
def process_observation(self, observation: Observation) -> str:
|
||
"""
|
||
Main entry point called by the agent each turn.
|
||
Generates reflection and pushes a compact line into the trajectory.
|
||
|
||
Args:
|
||
observation: The new game observation
|
||
|
||
Returns:
|
||
processed_observation: An updated observation with processed data
|
||
"""
|
||
|
||
"""
|
||
`-->` represents conversion performed by memory module
|
||
game_trajctory |-- [obs_i, action_i] |--> reflection
|
||
|
||
|
||
(inspired by LMAct)
|
||
Maybe we can add demonstrations as well
|
||
"""
|
||
game_state = observation.get_perception_summary()
|
||
|
||
prev_context = observation.game_trajectory.get() or ""
|
||
if observation.game_trajectory.background is None and observation.trajectory_includes_background:
|
||
observation.game_trajectory.set_background(observation.get_background() or "Background not available.")
|
||
|
||
if self.use_reflection:
|
||
reflection = self._reflect(
|
||
prev_context=prev_context,
|
||
current_state=str(game_state),
|
||
)
|
||
else:
|
||
reflection = None
|
||
|
||
observation = self.update_observation_memory(
|
||
observation=observation,
|
||
)
|
||
observation.reflection = reflection if self.use_reflection else None
|
||
|
||
return observation
|
||
|
||
def update_observation_memory(self, observation: Observation) -> str:
|
||
game_state = observation.get_perception_summary()
|
||
|
||
ts = datetime.datetime.now().isoformat(timespec="seconds")
|
||
game_state.pop("img_path")
|
||
|
||
if "processed_visual_description" in game_state and game_state["processed_visual_description"] is None:
|
||
game_state.pop("processed_visual_description")
|
||
|
||
# reflection excluded from game trajectory
|
||
# reflection will be extracted by the reasoning module
|
||
line = (
|
||
f"##Turn Hash\n[{ts}]\n"
|
||
f"###Obs\n{game_state}\n"
|
||
)
|
||
#f"###Reflection\n{reflection}\n"
|
||
|
||
# Get current trajectory content for summarization
|
||
current_trajectory = observation.game_trajectory.get() or ""
|
||
char_len = len(current_trajectory)
|
||
est_tokens = char_len // 3
|
||
|
||
# Check if we need to summarize before adding new entry
|
||
if self.use_summary and len(observation.game_trajectory.trajectory) >= self.max_memory or (est_tokens > 10_000 and "o3" in self.model_name):
|
||
# Trigger summarisation if > 10 000 tokens
|
||
|
||
# Only attempt summarization if we have substantial content
|
||
if len(current_trajectory.strip()) > 50: # Ensure we have meaningful content
|
||
print(f"[MemoryModule] Trajectory reached length: {len(observation.game_trajectory.trajectory)}.")
|
||
print(f"[MemoryModule] Current trajectory length: {char_len} chars")
|
||
print(f"[MemoryModule] Trajectory ≈ {est_tokens:,} tokens — summarizing…")
|
||
|
||
# Generate summary
|
||
new_summary = self._summarize(current_trajectory)
|
||
if new_summary and new_summary != "":
|
||
self.current_summary = new_summary
|
||
|
||
# Clear the trajectory and replace with summary
|
||
observation.game_trajectory.trajectory.clear()
|
||
|
||
# Clear the disk file as well since we're starting fresh
|
||
try:
|
||
with open(self.module_file, "w") as f:
|
||
json.dump([], f, indent=2)
|
||
except Exception as exc:
|
||
print(f"[MemoryModule] failed to clear log file: {exc}")
|
||
|
||
# Add summary as the first entry
|
||
summary_line = f"##TRAJECTORY SUMMARY\n{self.current_summary}\n\n"
|
||
observation.game_trajectory.add(summary_line)
|
||
|
||
# Persist summary to disk
|
||
self._append_to_log(summary_line)
|
||
|
||
print(f"[MemoryModule] Successfully generated and saved summary. Length: {len(self.current_summary)} chars")
|
||
else:
|
||
print(f"[MemoryModule] Failed to generate valid summary after all retries. Keeping existing trajectory.")
|
||
else:
|
||
print(f"[MemoryModule] Insufficient trajectory content ({len(current_trajectory)} chars) for summarization. Skipping.")
|
||
|
||
# add to dequeue
|
||
observation.game_trajectory.add(line)
|
||
# disk persistence
|
||
self._append_to_log(line)
|
||
|
||
return observation
|
||
|
||
def update_action_memory(self,
|
||
observation: Observation,
|
||
action: str | None,
|
||
thought: str | None) -> str:
|
||
"""
|
||
Main entry point called by the agent each turn.
|
||
Generates reflection and pushes a compact line into the trajectory.
|
||
|
||
Args:
|
||
observation: The new game observation
|
||
|
||
Returns:
|
||
processed_observation: An updated observation with processed data
|
||
"""
|
||
|
||
# build a single printable entry line
|
||
ts = datetime.datetime.now().isoformat(timespec="seconds")
|
||
line = (
|
||
f"###Action\n{action}\n"
|
||
f"###Thought\n{thought}\n"
|
||
)
|
||
|
||
# add to dequeue
|
||
observation.game_trajectory.add(line)
|
||
# disk persistence
|
||
self._append_to_log(line)
|
||
|
||
return observation
|
||
|
||
def get_memory_summary(self, observation) -> dict[str, str]:
|
||
"""
|
||
Provide the reasoning module with:
|
||
• up‑to‑N past lines (already formatted by GameTrajectory)
|
||
• no extra metadata dance
|
||
"""
|
||
past = observation.game_trajectory.get() or "No previous game states available."
|
||
latest = observation.game_trajectory.trajectory[-1] if observation.game_trajectory.trajectory else "N/A"
|
||
|
||
return {
|
||
"game_trajectory": past,
|
||
"current_state": latest, # includes (obs, action, thought)
|
||
"reflection": observation.reflection if hasattr(observation, 'reflection') and observation.reflection else "N/A",
|
||
}
|
||
|
||
def _parse_response(self, response):
|
||
"""
|
||
Parse the reflection response from the LLM.
|
||
|
||
Args:
|
||
response (str): The raw response from the LLM
|
||
|
||
Returns:
|
||
dict: Parsed reflection data
|
||
"""
|
||
|
||
if not response:
|
||
return {"reflection": "No reflection generated."}
|
||
|
||
# Try to extract reflection from structured format first
|
||
reflection_match = re.search(r'(?:^|\n)(?:#\s*)?reflection:(.+?)(?=(?:\n(?:#\s*)?[a-zA-Z]+:)|$)',
|
||
response, re.DOTALL | re.IGNORECASE)
|
||
|
||
if reflection_match:
|
||
# Extract the reflection content from the pattern match
|
||
reflection = reflection_match.group(1).strip()
|
||
else:
|
||
# If no structured format found, use the entire response
|
||
reflection = response.strip()
|
||
|
||
return {
|
||
"reflection": reflection
|
||
}
|