Spaces:
Running
Running
| """Short_term — per-agent, per-day operational memory. | |
| Stores the full day (plan, events, conversations, world snapshots, LLM | |
| daily summary) as one JSON file per persona per date, with atomic writes, | |
| and archives the day to Qdrant long-term memory at handoff. | |
| Architecture: the memory layer between the engine and long-term storage; | |
| written by world_engine.py, read by brain.py and day_planner.py. | |
| Design: short-term files are the operational record of the active day | |
| and are removed after successful archival. | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import json | |
| import os | |
| import re | |
| import tempfile | |
| import threading | |
| from datetime import datetime, timedelta, timezone | |
| from pathlib import Path | |
| from typing import Any, Optional | |
| from pydantic import BaseModel | |
| from src.config import DATA_DIR | |
| from src.core.log import get_logger | |
| logger = get_logger(__name__) | |
| # Directory for short-term memory files | |
| SHORT_TERM_DIR = DATA_DIR / "Short_term_db" | |
| _DAY_LOCKS_GUARD = threading.Lock() | |
| _DAY_LOCKS: dict[Path, threading.RLock] = {} | |
| # --------------------------------------------------------------------------- | |
| # Data Models (internal) | |
| # --------------------------------------------------------------------------- | |
| class DayEvent(BaseModel): | |
| """One event that happened during the day.""" | |
| timestamp: str # ISO format with timezone | |
| type: str # action_started, action_completed, conversation, world_snapshot, observation | |
| action: Optional[str] = None | |
| outcome: Optional[str] = None | |
| success: Optional[bool] = None | |
| location: Optional[str] = None | |
| with_agent: Optional[str] = None | |
| participants: Optional[list[str]] = None | |
| topic: Optional[str] = None | |
| summary: Optional[str] = None | |
| messages: Optional[list[dict]] = None | |
| details: Optional[dict] = None | |
| class ConversationEntry(BaseModel): | |
| """A conversation with another agent.""" | |
| timestamp: str | |
| participants: list[str] | |
| messages: list[dict] | |
| summary: str | |
| class WorldSnapshotEntry(BaseModel): | |
| """Periodic world state snapshot.""" | |
| timestamp: str | |
| location: str | |
| nearby_agents: list[str] | |
| weather: Optional[str] = None | |
| details: Optional[dict] = None | |
| class ShortTermDayData(BaseModel): | |
| """Complete day memory for one persona on one simulation date.""" | |
| persona_name: str | |
| date: str | |
| day_plan: list[dict] = [] | |
| events: list[DayEvent] = [] | |
| conversations: list[ConversationEntry] = [] | |
| world_snapshots: list[WorldSnapshotEntry] = [] | |
| daily_summary: Optional[str] = None | |
| archived_to_long_term: bool = False | |
| # --------------------------------------------------------------------------- | |
| # Helpers | |
| # --------------------------------------------------------------------------- | |
| def _persona_dir(persona_name: str) -> Path: | |
| """Get the directory for a persona's short-term files.""" | |
| safe_name = re.sub(r"[^A-Za-z0-9._-]+", "_", persona_name.strip()).strip("._-").lower() or "unknown" | |
| dir_path = SHORT_TERM_DIR / safe_name | |
| dir_path.mkdir(parents=True, exist_ok=True) | |
| return dir_path | |
| def _day_file_path(persona_name: str, date_str: str) -> Path: | |
| """Get the file path for a specific persona and date.""" | |
| return _persona_dir(persona_name) / f"{date_str}.json" | |
| def _day_lock(persona_name: str, date_str: str) -> threading.RLock: | |
| path = _day_file_path(persona_name, date_str).resolve() | |
| with _DAY_LOCKS_GUARD: | |
| return _DAY_LOCKS.setdefault(path, threading.RLock()) | |
| def _mutate_day_data(persona_name: str, date_str: str, mutate) -> None: | |
| """Serialize a complete same-day read/modify/write transaction.""" | |
| with _day_lock(persona_name, date_str): | |
| data = _load_day_data(persona_name, date_str) | |
| if data is None: | |
| data = ShortTermDayData(persona_name=persona_name, date=date_str) | |
| mutate(data) | |
| _save_day_data(persona_name, date_str, data) | |
| def date_from_simulation_time(time_str: str) -> str: | |
| """ | |
| Extract simulation date from a time string like '2026-07-03 06:00'. | |
| Returns 'YYYY-MM-DD'. | |
| """ | |
| # Expected format: "YYYY-MM-DD HH:MM" or just "YYYY-MM-DD HH:MM:SS" | |
| match = re.match(r"(\d{4}-\d{2}-\d{2})", time_str) | |
| if match: | |
| return match.group(1) | |
| # Fallback: try parsing | |
| try: | |
| dt = datetime.fromisoformat(time_str.replace("Z", "+00:00")) | |
| return dt.date().isoformat() | |
| except Exception: | |
| # Last resort: today's date (real) | |
| return datetime.now(timezone.utc).date().isoformat() | |
| def _load_day_data(persona_name: str, date_str: str) -> Optional[ShortTermDayData]: | |
| """Load day data from JSON file.""" | |
| path = _day_file_path(persona_name, date_str) | |
| if not path.exists(): | |
| return None | |
| try: | |
| raw = json.loads(path.read_text(encoding="utf-8")) | |
| return ShortTermDayData.model_validate(raw) | |
| except Exception as e: | |
| logger.error("Failed to load short-term data for %s on %s: %s", persona_name, date_str, e) | |
| return None | |
| def _save_day_data(persona_name: str, date_str: str, data: ShortTermDayData) -> None: | |
| """Atomically replace a day file so a process crash cannot leave partial JSON.""" | |
| path = _day_file_path(persona_name, date_str) | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| payload = data.model_dump_json(indent=2) | |
| fd, temporary_name = tempfile.mkstemp( | |
| prefix=f".{path.stem}.", suffix=".tmp", dir=path.parent, text=True, | |
| ) | |
| try: | |
| with os.fdopen(fd, "w", encoding="utf-8") as temporary: | |
| temporary.write(payload) | |
| temporary.flush() | |
| os.fsync(temporary.fileno()) | |
| os.replace(temporary_name, path) | |
| finally: | |
| if os.path.exists(temporary_name): | |
| os.unlink(temporary_name) | |
| logger.debug("Saved short-term data for %s on %s to %s", persona_name, date_str, path) | |
| # --------------------------------------------------------------------------- | |
| # Public API -- MemoryStreamProtocol (for tick_graph.py) | |
| # --------------------------------------------------------------------------- | |
| def add_memory(agent_id: str, content: str, importance: Optional[int] = None, date_str: Optional[str] = None) -> None: | |
| """ | |
| Add a memory entry for an agent. | |
| Stores as an event in the short-term file for the given (or current) date. | |
| """ | |
| if date_str is None: | |
| date_str = datetime.now(timezone.utc).date().isoformat() | |
| data = _load_day_data(agent_id, date_str) | |
| if data is None: | |
| data = ShortTermDayData(persona_name=agent_id, date=date_str) | |
| data.events.append(DayEvent( | |
| timestamp=datetime.now(timezone.utc).isoformat(), | |
| type="memory", | |
| details={"content": content, "importance": importance} | |
| )) | |
| _save_day_data(agent_id, date_str, data) | |
| def retrieve_memories(agent_id: str, query: str, k: int = 5, date_str: Optional[str] = None) -> list[str]: | |
| """Retrieve durable context through the Qdrant semantic-memory RAG layer.""" | |
| from src.agents.Long_term import get_retriever | |
| return get_retriever().retrieve(agent_id, query, k) | |
| # --------------------------------------------------------------------------- | |
| # Public API -- Day Planner Integration | |
| # --------------------------------------------------------------------------- | |
| def save_day_plan(persona_name: str, date_str: str, plan: list[dict]) -> None: | |
| """Save the generated day plan at the start of the day.""" | |
| data = _load_day_data(persona_name, date_str) | |
| if data is None: | |
| data = ShortTermDayData(persona_name=persona_name, date=date_str) | |
| data.day_plan = plan | |
| _save_day_data(persona_name, date_str, data) | |
| def load_day_plan(persona_name: str, date_str: str) -> list[dict]: | |
| """Load the day plan for a given date.""" | |
| data = _load_day_data(persona_name, date_str) | |
| return data.day_plan if data else [] | |
| # --------------------------------------------------------------------------- | |
| # Single-file hygiene (one short-term file per agent; clean start each run) | |
| # --------------------------------------------------------------------------- | |
| def is_pristine_plan(persona_name: str, date_str: str) -> bool: | |
| """True if the day file holds a plan but none of the run-time record yet | |
| (no events, conversations, or world snapshots).""" | |
| data = _load_day_data(persona_name, date_str) | |
| return bool( | |
| data and data.day_plan | |
| and not data.events and not data.conversations and not data.world_snapshots | |
| ) | |
| def reset_day_runtime(persona_name: str, date_str: str) -> None: | |
| """Clear the run-time record (events / conversations / snapshots) but keep | |
| the plan, so a day starts fresh — the file holds only the plan again.""" | |
| data = _load_day_data(persona_name, date_str) | |
| if data is None: | |
| return | |
| data.events = [] | |
| data.conversations = [] | |
| data.world_snapshots = [] | |
| data.archived_to_long_term = False | |
| _save_day_data(persona_name, date_str, data) | |
| def list_day_files(persona_name: str) -> list[str]: | |
| """All simulation dates that currently have a short-term file for this agent.""" | |
| d = _persona_dir(persona_name) | |
| return sorted(p.stem for p in d.glob("*.json")) | |
| def archive_copy_no_llm(persona_name: str, date_str: str) -> bool: | |
| """Index a completed operational day without making a summary LLM call. | |
| This is used only to clean up stray short-term files. It never writes a | |
| local long-term JSON archive, and it retains the source file on failure. | |
| """ | |
| data = _load_day_data(persona_name, date_str) | |
| if data is None: | |
| return False | |
| try: | |
| from src.agents.Long_term import get_retriever | |
| return bool(get_retriever().index_archive(persona_name, data.model_dump())) | |
| except Exception as exc: | |
| logger.warning("[Short_term] unable to index stray day %s for %s: %s", date_str, persona_name, exc) | |
| return False | |
| def consolidate_to_single_day(persona_name: str, keep_date: str) -> list[str]: | |
| """Ensure the agent keeps only ONE short-term file (for `keep_date`). | |
| Any other day files are indexed into Qdrant (no LLM) and removed only | |
| after Qdrant confirms the write. Returns the dates safely removed. | |
| """ | |
| moved: list[str] = [] | |
| for date_str in list_day_files(persona_name): | |
| if date_str == keep_date: | |
| continue | |
| if archive_copy_no_llm(persona_name, date_str): | |
| clear_short_term_data(persona_name, date_str) | |
| moved.append(date_str) | |
| if moved: | |
| logger.info( | |
| "[Short_term] consolidated %s: indexed %d stray day file(s) into Qdrant (%s)", | |
| persona_name, len(moved), ", ".join(moved), | |
| ) | |
| return moved | |
| def append_event(persona_name: str, date_str: str, event: dict) -> None: | |
| """ | |
| Append an event during the day. | |
| event dict should have: type, action, outcome, success, location, details, etc. | |
| """ | |
| entry = dict(event) | |
| entry.setdefault("timestamp", datetime.now(timezone.utc).isoformat()) | |
| _mutate_day_data( | |
| persona_name, date_str, | |
| lambda data: data.events.append(DayEvent(**entry)), | |
| ) | |
| def append_conversation(persona_name: str, date_str: str, conversation: dict) -> None: | |
| """ | |
| Append a conversation during the day. | |
| conversation dict: participants, messages, summary | |
| """ | |
| entry = dict(conversation) | |
| entry.setdefault("timestamp", datetime.now(timezone.utc).isoformat()) | |
| _mutate_day_data( | |
| persona_name, date_str, | |
| lambda data: data.conversations.append(ConversationEntry(**entry)), | |
| ) | |
| def append_world_snapshot(persona_name: str, date_str: str, snapshot: dict) -> None: | |
| """ | |
| Append a world state snapshot during the day. | |
| snapshot dict: location, nearby_agents, weather, details | |
| """ | |
| entry = dict(snapshot) | |
| entry.setdefault("timestamp", datetime.now(timezone.utc).isoformat()) | |
| _mutate_day_data( | |
| persona_name, date_str, | |
| lambda data: data.world_snapshots.append(WorldSnapshotEntry(**entry)), | |
| ) | |
| def get_yesterday_data(persona_name: str, today_date: str) -> Optional[ShortTermDayData]: | |
| """ | |
| Get the full yesterday's data for a persona. | |
| today_date: 'YYYY-MM-DD' (simulation date) | |
| Returns yesterday's ShortTermDayData or None. | |
| """ | |
| try: | |
| today_dt = datetime.fromisoformat(today_date).date() | |
| yesterday = (today_dt - timedelta(days=1)).isoformat() | |
| except Exception: | |
| # Fallback to real date | |
| yesterday = (datetime.now(timezone.utc).date() - timedelta(days=1)).isoformat() | |
| return _load_day_data(persona_name, yesterday) | |
| def get_yesterday_summary(persona_name: str, today_date: str) -> Optional[str]: | |
| """ | |
| Get a text summary of yesterday for the day planner prompt. | |
| Returns None if no yesterday data exists. | |
| """ | |
| from src.agents.Long_term import get_retriever | |
| return get_retriever().rolling_summary(persona_name, days=1, before_date=today_date) | |
| def get_relevant_memories(persona_name: str, today_date: str, query: str, k: int = 10) -> list[str]: | |
| """ | |
| Get relevant long-term memories for model prompts through semantic search. | |
| """ | |
| return retrieve_memories(persona_name, query, k, date_str=today_date) | |
| # --------------------------------------------------------------------------- | |
| # Public API -- End of Day Lifecycle | |
| # --------------------------------------------------------------------------- | |
| async def generate_daily_summary(persona_name: str, date_str: str) -> str: | |
| """ | |
| Generate an LLM summary of the day's events for long-term storage. | |
| Uses the same gemini_client as day_planner. | |
| """ | |
| from src.llm.gemini_client import call_gemini | |
| from pydantic import BaseModel | |
| from typing import List | |
| class SummaryOutput(BaseModel): | |
| summary: str | |
| key_events: List[str] | |
| data = _load_day_data(persona_name, date_str) | |
| if data is None or not data.events: | |
| return f"No events recorded for {persona_name} on {date_str}." | |
| # Build event text | |
| event_lines = [] | |
| for e in data.events: | |
| if e.type == "action_completed": | |
| event_lines.append(f"- {e.action} at {e.location or 'unknown'} (outcome: {e.outcome}, success: {e.success})") | |
| elif e.type == "conversation": | |
| event_lines.append(f"- Conversation with {e.with_agent or ', '.join(e.participants) if e.participants else 'unknown'}: {e.summary or e.topic}") | |
| elif e.type == "observation": | |
| event_lines.append(f"- Observed: {e.details}") | |
| event_text = "\n".join(event_lines) | |
| system_prompt = ( | |
| "You are summarizing one day in the life of a simulated college student agent. " | |
| "You will be provided the Initial day plan, the conversations the agent had throughout and some extra stuff. " | |
| "Write a concise summary of about 100 words prioritizing the most important:\n" | |
| "- High-priority tasks and whether they were completed\n" | |
| "- Important events and social interactions\n" | |
| "- Significant conversations and decisions\n" | |
| "- Anything likely to matter in future days\n" | |
| "Only focus on routine low-value chores unless no important event happened that day.\n" | |
| "Also list 3-5 key events as bullet points." | |
| ) | |
| user_prompt = ( | |
| f"Persona: {persona_name}\n" | |
| f"Date: {date_str}\n" | |
| f"Day plan: {len(data.day_plan)} scheduled actions\n\n" | |
| f"Initial day plan:\n" + | |
| "\n".join(f" - {a.get('start','?')}-{a.get('end','?')} {a.get('action','?')} at {a.get('location_id','?')}" for a in data.day_plan[:15]) + | |
| f"\n\nEvents:\n{event_text}\n\n" | |
| f"Conversations ({len(data.conversations)}):\n" + | |
| "\n".join(f"- {c.summary} (with {', '.join(c.participants)})" for c in data.conversations[:5]) + | |
| "\n\nGenerate the daily summary (~100 words) and key events list." | |
| ) | |
| try: | |
| from src.config import SUMMARY_TEMPERATURE | |
| result = call_gemini(system_prompt, user_prompt, SummaryOutput, "default", temperature=SUMMARY_TEMPERATURE) | |
| # Format for long-term storage | |
| key_events_text = "\n".join(f"- {evt}" for evt in result.key_events) | |
| return f"Summary: {result.summary}\nKey events:\n{key_events_text}" | |
| except Exception as e: | |
| logger.error("Failed to generate LLM daily summary for %s on %s: %s", persona_name, date_str, e) | |
| # Fallback: programmatic summary | |
| parts = [f"{persona_name} on {date_str}:"] | |
| for e in data.events: | |
| if e.type == "action_completed": | |
| parts.append(f" - {e.action} ({e.outcome}, success={e.success})") | |
| return "\n".join(parts) | |
| async def finalize_day(persona_name: str, date_str: str) -> dict: | |
| """ | |
| End-of-day processing: | |
| 1. Generate LLM summary of the day | |
| 2. Mark as archived | |
| 3. Return data ready for long-term storage | |
| """ | |
| data = _load_day_data(persona_name, date_str) | |
| if data is None: | |
| return {"summary": f"No data for {persona_name} on {date_str}", "full_data": None} | |
| # Generate summary | |
| daily_summary = await generate_daily_summary(persona_name, date_str) | |
| data.daily_summary = daily_summary | |
| data.archived_to_long_term = True | |
| # Save final version with summary | |
| _save_day_data(persona_name, date_str, data) | |
| # Return data for long-term storage | |
| return { | |
| "summary": daily_summary, | |
| "full_data": data.model_dump() | |
| } | |
| async def archive_to_long_term(persona_name: str, date_str: str) -> dict: | |
| """ | |
| End-of-day archival: generate a daily summary and persist its durable | |
| records directly to Qdrant. No Long_term_db JSON file is created. | |
| Returns the archive payload dict with keys: summary, persona_name, date, | |
| event_count, conversation_count. | |
| """ | |
| result = await finalize_day(persona_name, date_str) | |
| data = result.get("full_data") | |
| if data is None: | |
| return {"summary": f"No data for {persona_name} on {date_str}"} | |
| # Offline/local runs intentionally disable the Qdrant-backed long-term | |
| # store. Treat this as an explicit, successful skip rather than trying a | |
| # known-unavailable retriever and turning every midnight handoff into an | |
| # archive failure. | |
| from src import config as _cfg | |
| if not _cfg.SEMANTIC_MEMORY_ENABLED: | |
| logger.info("[Short_term] semantic archival disabled; skipped %s on %s", persona_name, date_str) | |
| return { | |
| "summary": result.get("summary", ""), | |
| "persona_name": persona_name, | |
| "date": date_str, | |
| "event_count": len(data.get("events", [])), | |
| "conversation_count": len(data.get("conversations", [])), | |
| "skipped": True, | |
| } | |
| try: | |
| from src.agents.Long_term import get_retriever | |
| retriever = get_retriever() | |
| indexed = await asyncio.to_thread(retriever.index_archive, persona_name, data) | |
| if indexed <= 0 and data.get("daily_summary"): | |
| raise RuntimeError("Qdrant did not acknowledge any durable memory records") | |
| except Exception as exc: | |
| # The short-term source remains in place so an operator can retry. Do | |
| # not silently claim archival success when Qdrant is the only durable | |
| # long-term store. | |
| logger.error("[Short_term] Qdrant archive failed for %s: %s", persona_name, exc) | |
| raise | |
| logger.info( | |
| "[Short_term] archived %s on %s into Qdrant (%d events, %d conversations)", | |
| persona_name, date_str, | |
| len(data.get("events", [])), len(data.get("conversations", [])), | |
| ) | |
| return { | |
| "summary": result.get("summary", ""), | |
| "persona_name": persona_name, | |
| "date": date_str, | |
| "event_count": len(data.get("events", [])), | |
| "conversation_count": len(data.get("conversations", [])), | |
| } | |
| def clear_short_term_data(persona_name: str, date_str: str) -> bool: | |
| """Delete the short-term file after successful long-term archival. | |
| Not called automatically — archival keeps the short-term copy in place so | |
| the running day is not disturbed. Exposed for maintenance/cleanup use. | |
| """ | |
| path = _day_file_path(persona_name, date_str) | |
| if path.exists(): | |
| path.unlink() | |
| logger.info("Cleared short-term data for %s on %s", persona_name, date_str) | |
| return True | |
| return False | |
| # --------------------------------------------------------------------------- | |
| # MemoryStreamProtocol adapter | |
| # --------------------------------------------------------------------------- | |
| class ShortTermMemoryStream: | |
| """ | |
| Adapter that wraps Short_term.py's module-level functions as | |
| MemoryStreamProtocol (defined in tick_graph.py) with sim-date awareness. | |
| Usage: | |
| stream = ShortTermMemoryStream("2026-07-03") | |
| # Tick graph nodes call: | |
| stream.add_memory(agent_id, content, importance) | |
| stream.retrieve_memories(agent_id, query, k) | |
| """ | |
| def __init__(self, initial_date: str = "") -> None: | |
| self._date_str: str = initial_date or datetime.now(timezone.utc).date().isoformat() | |
| def date_str(self) -> str: | |
| return self._date_str | |
| def set_date(self, date_str: str) -> None: | |
| """Update the simulation date (called before each tick).""" | |
| self._date_str = date_str | |
| def add_memory(self, agent_id: str, content: str, importance: Optional[int] = None) -> None: | |
| add_memory(agent_id, content, importance, date_str=self._date_str) | |
| def retrieve_memories(self, agent_id: str, query: str, k: int = 5) -> list[str]: | |
| return retrieve_memories(agent_id, query, k, date_str=self._date_str) | |