import os
import sqlite3
import logging
import uuid
from contextlib import contextmanager
from pathlib import Path
from typing import List, Dict, Any, Optional, Union
from datetime import datetime
from langchain_core.messages import BaseMessage, HumanMessage
log = logging.getLogger(__name__)
# Constants
LILITH_HOME = Path(os.getenv("LILITH_HOME", ".lilith"))
MEMORY_DB_PATH = LILITH_HOME / "long_term_memory.sqlite"
MEMORY_CONTEXT_CHAR_BUDGET = 3000
MIN_MESSAGES_FOR_EXTRACTION = 2
MEMORY_LOG_PREVIEW_CHARS = 300
def _content_to_text(content: Any) -> str:
if content is None:
return ""
if isinstance(content, str):
return content
if isinstance(content, list):
parts = []
for item in content:
if isinstance(item, str):
parts.append(item)
elif isinstance(item, dict):
text = item.get("text")
if text is not None:
parts.append(str(text))
elif "content" in item:
nested = _content_to_text(item["content"])
if nested:
parts.append(nested)
else:
parts.append(str(item))
return "\n".join(part for part in parts if part)
return str(content)
def _memory_content_to_text(content: Any) -> str:
if hasattr(content, "content"):
return _content_to_text(content.content)
return _content_to_text(content)
def _preview_for_log(text: str, limit: int = MEMORY_LOG_PREVIEW_CHARS) -> str:
compact = " ".join(str(text).split())
if len(compact) <= limit:
return compact
return compact[: limit - 1] + "…"
def _print_memory_log(message: str, *args: Any) -> None:
text = message % args if args else message
print(text, flush=True)
class MemoryStore:
def __init__(self, db_path: Union[Path, str] = MEMORY_DB_PATH):
self._in_memory = (str(db_path) == ":memory:")
self.db_path = db_path if self._in_memory else Path(db_path)
self._mem_conn: Optional[sqlite3.Connection] = None
self._init_db()
def _connect(self) -> sqlite3.Connection:
if self._in_memory:
if self._mem_conn is None:
self._mem_conn = sqlite3.connect(":memory:", check_same_thread=False)
return self._mem_conn
return sqlite3.connect(str(self.db_path))
def close(self):
if self._mem_conn is not None:
self._mem_conn.close()
self._mem_conn = None
def _init_db(self):
if not self._in_memory:
Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
conn = self._connect()
with conn:
conn.execute("""
CREATE TABLE IF NOT EXISTS memories (
id TEXT PRIMARY KEY,
content TEXT NOT NULL,
type TEXT DEFAULT 'fact',
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
)
""")
conn.execute("""
CREATE TABLE IF NOT EXISTS episodes (
id TEXT PRIMARY KEY,
task TEXT NOT NULL,
summary TEXT NOT NULL,
outcome TEXT,
created_at TEXT NOT NULL
)
""")
if not self._in_memory:
conn.close()
def get_all_memories(self) -> List[Dict[str, Any]]:
conn = self._connect()
conn.row_factory = sqlite3.Row
cur = conn.cursor()
cur.execute("SELECT * FROM memories ORDER BY updated_at DESC")
rows = [dict(r) for r in cur.fetchall()]
if not self._in_memory:
conn.close()
return rows
def save_memories(self, memories: List[Dict[str, Any]], allow_empty: bool = False):
"""Replaces the memories table with the provided list (active compression)."""
if not memories:
existing = self.get_all_memories()
if existing and not allow_empty:
_print_memory_log("[memory] save_memories called with empty list while %d facts exist — refusing to wipe",
len(existing))
log.warning("[memory] save_memories called with empty list while %d facts exist — refusing to wipe",
len(existing))
return
conn = self._connect()
with conn:
conn.execute("DELETE FROM memories")
for m in memories:
now = datetime.now().isoformat()
created_at = m.get("created_at", now)
updated_at = m.get("updated_at", now)
conn.execute(
"INSERT INTO memories (id, content, type, created_at, updated_at) VALUES (?, ?, ?, ?, ?)",
(m.get("id", str(uuid.uuid4())), m["content"], m.get("type", "fact"), created_at, updated_at)
)
if not self._in_memory:
conn.close()
def delete_memory_prefix(self, prefix: str) -> int:
if not prefix:
return 0
matching_ids = [
row["id"]
for row in self.get_all_memories()
if row["id"].startswith(prefix)
]
if not matching_ids:
return 0
conn = self._connect()
with conn:
for memory_id in matching_ids:
conn.execute("DELETE FROM memories WHERE id = ?", (memory_id,))
if not self._in_memory:
conn.close()
return len(matching_ids)
def add_episode(self, task: str, summary: str, outcome: str):
conn = self._connect()
with conn:
now = datetime.now().isoformat()
conn.execute(
"INSERT INTO episodes (id, task, summary, outcome, created_at) VALUES (?, ?, ?, ?, ?)",
(str(uuid.uuid4()), task, summary, outcome, now)
)
if not self._in_memory:
conn.close()
def get_recent_episodes(self, limit: int = 3) -> List[Dict[str, Any]]:
conn = self._connect()
conn.row_factory = sqlite3.Row
cur = conn.cursor()
cur.execute("SELECT * FROM episodes ORDER BY created_at DESC LIMIT ?", (limit,))
rows = [dict(r) for r in cur.fetchall()]
if not self._in_memory:
conn.close()
return rows
_store = MemoryStore()
def _set_store(store: MemoryStore) -> MemoryStore:
"""Swap the module-level store and return the previous one."""
global _store
prev, _store = _store, store
return prev
@contextmanager
def ephemeral_memory(db_path: Union[str, Path] = ":memory:"):
"""Context manager that replaces the global _store with a fresh, isolated
MemoryStore for the duration of the block. On exit the store is closed and
the previous store is restored. Use db_path=':memory:' (default) for tests
or GAIA benchmarking where cross-contamination must be prevented.
Example::
with ephemeral_memory():
result = graph.invoke(state, config)
# any memory writes here are discarded on exit
"""
fresh = MemoryStore(db_path)
prev = _set_store(fresh)
try:
yield fresh
finally:
_set_store(prev)
fresh.close()
def _is_remove_doc(content: Any) -> bool:
return hasattr(content, "__repr_name__") and content.__repr_name__() == "RemoveDoc"
def extract_and_compress_facts(messages: List[BaseMessage], model) -> None:
"""
Extracts new facts and merges them with existing ones using langmem's manager.
Implements professional reflection and conflict resolution.
"""
from langmem import create_memory_manager
from langmem.knowledge.extraction import Memory
_print_memory_log("[memory] Running langmem memory management...")
log.info("[memory] Running langmem memory management...")
try:
# 1. Get existing memories from our local store
existing_rows = _store.get_all_memories()
existing_memories = [(m["id"], Memory(content=m["content"])) for m in existing_rows]
existing_by_content = {m["content"]: m for m in existing_rows}
existing_by_id = {m["id"]: m for m in existing_rows}
# 2. Initialize langmem manager (it's a Runnable)
# We use the default schema which is basically content strings
manager = create_memory_manager(model, enable_deletes=True)
# 3. Invoke manager with history and existing knowledge
# The manager will return the full updated list of memories
result = manager.invoke({
"messages": messages,
"existing": existing_memories
})
# 4. Save results back to our local persistent SQLite
if isinstance(result, list):
# The result is a list of ExtractedMemory objects or tuples depending on version
# Usually it's (id, content) or just objects. We'll be robust here.
updated_facts = []
removed_ids = set()
for item in result:
item_id = getattr(item, "id", None)
content = getattr(item, "content", None)
if isinstance(item, tuple):
if len(item) > 0:
item_id = item[0]
if len(item) > 1 and content is None:
content = item[1]
elif isinstance(item, dict):
item_id = item.get("id", item_id)
content = item.get("content", content)
stable_id = str(item_id) if item_id else None
if _is_remove_doc(content):
if stable_id:
removed_ids.add(stable_id)
continue
if content:
text = _memory_content_to_text(content)
existing = existing_by_id.get(stable_id) if stable_id else existing_by_content.get(text)
fact = {"content": text}
if stable_id:
fact["id"] = stable_id
elif existing:
fact["id"] = existing["id"]
if existing:
fact["created_at"] = existing["created_at"]
if existing["content"] == text:
fact["updated_at"] = existing["updated_at"]
if fact.get("id") not in removed_ids:
updated_facts.append(fact)
if updated_facts or removed_ids:
_store.save_memories(updated_facts, allow_empty=bool(removed_ids))
_print_memory_log("[memory] langmem updated store to %d facts.", len(updated_facts))
log.info("[memory] langmem updated store to %d facts.", len(updated_facts))
else:
_print_memory_log("[memory] langmem returned empty result — keeping existing facts")
log.info("[memory] langmem returned empty result — keeping existing facts")
summarize_episode(messages, model)
except Exception as exc:
_print_memory_log("[memory] langmem extraction failed: %s: %s", type(exc).__name__, exc)
log.exception("[memory] langmem extraction failed")
# Fallback to summarize episode if manager fails
summarize_episode(messages, model)
def summarize_episode(messages: List[BaseMessage], model) -> None:
"""Summarizes the trajectory to help avoid future mistakes."""
_print_memory_log("[memory] Summarizing task episode...")
log.info("[memory] Summarizing task episode...")
try:
initial_question = ""
outcome = "success"
tool_names_used = set()
fetch_url_used = False
search_count = 0
for m in messages:
content = _content_to_text(m.content)
if isinstance(m, HumanMessage) and not initial_question:
initial_question = content
if "ERROR" in content.upper():
outcome = "failed/struggled"
if hasattr(m, "tool_calls"):
for tc in (m.tool_calls or []):
name = tc.get("name", "") if isinstance(tc, dict) else ""
tool_names_used.add(name)
if name == "fetch_url":
fetch_url_used = True
if name == "web_search":
search_count += 1
quality_notes = []
if search_count > 0 and not fetch_url_used:
quality_notes.append("Agent searched but never read any source pages (fetch_url unused)")
if search_count > 8:
quality_notes.append(f"High search volume ({search_count} searches) — possible over-searching")
quality_section = ""
if quality_notes:
quality_section = "\nQuality flags: " + "; ".join(quality_notes)
prompt = f"""
Summarize this task trajectory for Lilith's 'Episodic Memory'.
Initial Question: {initial_question}
Outcome: {outcome}
Tools used: {', '.join(sorted(tool_names_used)) if tool_names_used else 'none'}
{quality_section}
Briefly explain:
1. What was the goal?
2. What tools worked? What failed?
3. What is the 'lesson learned' for next time?
4. If quality flags are present, note them as areas for improvement.
Keep it under 150 words.
"""
response = model.invoke(prompt)
summary = _content_to_text(response.content)
_store.add_episode(initial_question, summary, outcome)
_print_memory_log(
"[memory] Episode saved: task=%r outcome=%r summary=%r",
_preview_for_log(initial_question),
outcome,
_preview_for_log(summary),
)
log.info(
"[memory] Episode saved: task=%r outcome=%r summary=%r",
_preview_for_log(initial_question),
outcome,
_preview_for_log(summary),
)
except Exception as exc:
_print_memory_log("[memory] Summarization failed: %s: %s", type(exc).__name__, exc)
log.exception("[memory] Summarization failed")
def _relevance_score(text: str, query_tokens: set) -> float:
"""Simple word-overlap score between a fact and the query (Jaccard-like)."""
if not query_tokens:
return 0.0
fact_tokens = set(text.lower().split())
return len(fact_tokens & query_tokens) / (len(fact_tokens | query_tokens) or 1)
def retrieve_relevant_context(query: str, char_budget: int = MEMORY_CONTEXT_CHAR_BUDGET) -> str:
"""Fetches facts and recent episodes ranked by relevance, capped by char_budget.
If truncated, appends a note instructing the agent to call search_memory for more."""
try:
query_tokens = set(query.lower().split())
facts = _store.get_all_memories()
episodes = _store.get_recent_episodes(limit=5)
facts.sort(key=lambda m: _relevance_score(m["content"], query_tokens), reverse=True)
context_parts = []
budget_remaining = char_budget
if facts:
included, total = [], len(facts)
for m in facts:
line = f"- {m['content']}"
if budget_remaining - len(line) - 1 > 0:
included.append(line)
budget_remaining -= len(line) + 1
else:
break
fact_block = "\n" + "\n".join(included) + "\n"
omitted = total - len(included)
if omitted:
fact_block += f"\n"
context_parts.append(fact_block)
if episodes and budget_remaining > 0:
included_ep = []
for e in episodes:
line = f"Task: {e['task']}\nSummary: {e['summary']}"
if budget_remaining - len(line) - 2 > 0:
included_ep.append(line)
budget_remaining -= len(line) + 2
else:
break
if included_ep:
ep_block = "\n" + "\n\n".join(included_ep) + "\n"
omitted_ep = len(episodes) - len(included_ep)
if omitted_ep:
ep_block += f"\n"
context_parts.append(ep_block)
return "\n\n".join(context_parts)
except Exception as exc:
_print_memory_log("[memory] Retrieval failed: %s: %s", type(exc).__name__, exc)
log.exception("[memory] Retrieval failed")
return ""
def search_memory_store(query: str, max_results: int = 10) -> str:
"""Keyword search across all facts and episodes. Called by the agent when
the system-prompt injection was truncated or the query needs deeper lookup."""
try:
query_tokens = set(query.lower().split())
facts = _store.get_all_memories()
episodes = _store.get_recent_episodes(limit=50)
scored_facts = sorted(
[(f, _relevance_score(f["content"], query_tokens)) for f in facts],
key=lambda x: x[1], reverse=True
)
scored_eps = sorted(
[(e, _relevance_score(e["task"] + " " + e["summary"], query_tokens)) for e in episodes],
key=lambda x: x[1], reverse=True
)
results = []
for m, score in scored_facts[:max_results]:
if score > 0:
results.append((score, f"[fact] {m['content']}"))
for e, score in scored_eps[:max_results]:
if score > 0:
results.append((score, f"[episode] Task: {e['task']}\n Summary: {e['summary']}"))
results.sort(key=lambda x: x[0], reverse=True)
parts = [text for _, text in results[:max_results]]
if not parts:
return "No matching memories found."
return "\n\n".join(parts)
except Exception as exc:
_print_memory_log("[memory] search_memory_store failed: %s: %s", type(exc).__name__, exc)
log.exception("[memory] search_memory_store failed")
return "Memory search failed."