mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-08 08:11:07 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
43724ab962 | ||
|
|
926ba86143 | ||
|
|
373b40e605 | ||
|
|
4aa310a214 | ||
|
|
e4cab9b7db |
No files matched your search
@@ -370,9 +370,9 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
self._index_session(session)
|
||||
return ok
|
||||
|
||||
def _save_imported_session(self, session: SessionData) -> bool:
|
||||
def _save_imported_session(self, session: SessionData, *, overwrite: bool = True) -> bool:
|
||||
"""Refresh recall and routing after a verbatim portable restore."""
|
||||
ok = super()._save_imported_session(session)
|
||||
ok = super()._save_imported_session(session, overwrite=overwrite)
|
||||
if ok:
|
||||
try:
|
||||
filepath = self._get_session_path(session.session_id)
|
||||
|
||||
@@ -216,7 +216,7 @@ class SqliteTranscriptStore(DefaultSessionStore):
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return None
|
||||
|
||||
def _write_row(self, session: SessionData, conn=None) -> bool:
|
||||
def _write_row(self, session: SessionData, conn=None, *, overwrite: bool = True) -> bool:
|
||||
own_lock = conn is None
|
||||
if conn is None:
|
||||
conn = self._connect()
|
||||
@@ -235,17 +235,21 @@ class SqliteTranscriptStore(DefaultSessionStore):
|
||||
getattr(session, "updated_at", None),
|
||||
)
|
||||
sql = (
|
||||
"INSERT OR REPLACE INTO sessions "
|
||||
("INSERT OR REPLACE" if overwrite else "INSERT OR IGNORE") + " INTO sessions "
|
||||
"(session_id, data, agent_name, gateway_session_id, agent_id, updated_at) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?)"
|
||||
)
|
||||
try:
|
||||
if own_lock:
|
||||
with self._db_lock:
|
||||
conn.execute(sql, params)
|
||||
cursor = conn.execute(sql, params)
|
||||
else:
|
||||
conn.execute(sql, params)
|
||||
cursor = conn.execute(sql, params)
|
||||
if not overwrite and cursor.rowcount == 0:
|
||||
raise FileExistsError(session.session_id)
|
||||
return True
|
||||
except FileExistsError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.error("Failed to write session %s: %s", session.session_id, exc)
|
||||
return False
|
||||
@@ -288,6 +292,52 @@ class SqliteTranscriptStore(DefaultSessionStore):
|
||||
self._cache[session.session_id] = session
|
||||
return True
|
||||
|
||||
def _save_imported_session(self, session: SessionData, *, overwrite: bool = True) -> bool:
|
||||
"""Restore to the database without applying the destination's window."""
|
||||
session.updated_at = datetime.now(timezone.utc).isoformat()
|
||||
with self._db_lock:
|
||||
# INSERT OR IGNORE makes no-clobber atomic across store instances;
|
||||
# checking session_exists before a replace would still race.
|
||||
if not self._write_row(session, overwrite=overwrite):
|
||||
return False
|
||||
with self._lock:
|
||||
self._cache[session.session_id] = session
|
||||
return True
|
||||
|
||||
def export_all(self) -> Dict[str, Any]:
|
||||
"""Export durable database rows, not legacy JSON sidecar files."""
|
||||
return {
|
||||
"version": self.PORTABLE_VERSION,
|
||||
"sessions": [data for data in self._all_rows() if isinstance(data, dict)],
|
||||
}
|
||||
|
||||
def _collect_lineage(
|
||||
self, session: SessionData, *, exclude: str
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Collect persisted database records from the same conversation chain."""
|
||||
lineage = self._lineage_key(session.to_dict())
|
||||
if not lineage:
|
||||
return []
|
||||
|
||||
def record_lineage(raw):
|
||||
# Preserve the writer's Python JSON semantics (including NaN)
|
||||
# and the shared lineage precedence, without retaining other rows.
|
||||
try:
|
||||
data = json.loads(raw)
|
||||
return self._lineage_key(data) if isinstance(data, dict) else None
|
||||
except (json.JSONDecodeError, TypeError, AttributeError):
|
||||
return None
|
||||
|
||||
conn = self._connect()
|
||||
with self._db_lock:
|
||||
conn.create_function("portable_session_lineage", 1, record_lineage)
|
||||
rows = conn.execute(
|
||||
"SELECT data FROM sessions WHERE session_id != ? "
|
||||
"AND portable_session_lineage(data) = ? ORDER BY updated_at DESC",
|
||||
(exclude, lineage),
|
||||
).fetchall()
|
||||
return [json.loads(row[0]) for row in rows]
|
||||
|
||||
def _modify_session_locked(
|
||||
self,
|
||||
session_id: str,
|
||||
|
||||
@@ -111,8 +111,8 @@ def test_write_between_import_save_and_index_keeps_newer_content(make_store, mon
|
||||
resume = threading.Event()
|
||||
original = DefaultSessionStore._save_imported_session
|
||||
|
||||
def delayed_save(store, session):
|
||||
result = original(store, session)
|
||||
def delayed_save(store, session, **kwargs):
|
||||
result = original(store, session, **kwargs)
|
||||
saved.set()
|
||||
assert resume.wait(5), "concurrent writer did not finish"
|
||||
return result
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
"""Portable session operations use the SQLite transcript persistence backend."""
|
||||
|
||||
import json
|
||||
import tracemalloc
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
from praisonaiagents.session.sqlite_transcript_store import SqliteTranscriptStore
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stores(tmp_path):
|
||||
opened = []
|
||||
|
||||
def make(name, backend=SqliteTranscriptStore, **kwargs):
|
||||
store = backend(session_dir=str(tmp_path / name), **kwargs)
|
||||
opened.append(store)
|
||||
return store
|
||||
|
||||
yield make
|
||||
for store in opened:
|
||||
conn = getattr(store, "_conn", None)
|
||||
if conn is not None:
|
||||
conn.close()
|
||||
|
||||
|
||||
def test_export_all_reads_database_not_json_sidecars(stores):
|
||||
store = stores("source")
|
||||
assert store.add_message("one", "user", "First")
|
||||
assert store.add_message("two", "assistant", "Second")
|
||||
from pathlib import Path
|
||||
(Path(store.session_dir) / "unrelated.json").write_text(
|
||||
json.dumps({"session_id": "sidecar", "messages": []}), encoding="utf-8"
|
||||
)
|
||||
payload = store.export_all()
|
||||
assert payload["version"] == store.PORTABLE_VERSION
|
||||
assert {record["session_id"] for record in payload["sessions"]} == {"one", "two"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key", ["lineage_id", "root_session_id", "thread_id"])
|
||||
def test_single_export_includes_database_lineage(stores, key):
|
||||
store = stores("source")
|
||||
for session_id in ("first", "continuation", "unrelated"):
|
||||
assert store.add_message(session_id, "user", session_id)
|
||||
assert store.update_session_metadata(session_id, **{key: "shared" if session_id != "unrelated" else "other"})
|
||||
payload = store.export_session("continuation")
|
||||
assert {record["session_id"] for record in payload["sessions"]} == {"first", "continuation"}
|
||||
assert [record["session_id"] for record in store.export_session("continuation", include_lineage=False)["sessions"]] == ["continuation"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source_backend", [DefaultSessionStore, SqliteTranscriptStore])
|
||||
@pytest.mark.parametrize("target_backend", [DefaultSessionStore, SqliteTranscriptStore])
|
||||
def test_cross_backend_restore_survives_reopen_without_window_truncation(stores, source_backend, target_backend):
|
||||
source = stores("source", source_backend)
|
||||
for index in range(6):
|
||||
assert source.add_message("reference", "user", f"message-{index}")
|
||||
assert source.set_gateway_info("reference", gateway_session_id="live", agent_id="agent")
|
||||
payload = source.export_all()
|
||||
assert len(payload["sessions"]) == 1
|
||||
target = stores("target", target_backend, active_window=2, retention="truncate")
|
||||
# An existing SQLite row prevents legacy JSON migration from masking a
|
||||
# restore incorrectly written to a sidecar file instead of the database.
|
||||
assert target.add_message("existing", "user", "Existing history")
|
||||
assert target.import_sessions(payload).imported == 1
|
||||
assert target.session_exists("reference")
|
||||
target.invalidate_cache()
|
||||
reopened = stores("target", target_backend, active_window=2, retention="truncate")
|
||||
for current in (target, reopened):
|
||||
assert [message["content"] for message in current.get_chat_history("reference")] == [f"message-{index}" for index in range(6)]
|
||||
session = current.get_session("reference")
|
||||
assert session.gateway_session_id is None
|
||||
assert session.agent_id is None
|
||||
if target_backend is SqliteTranscriptStore:
|
||||
from pathlib import Path
|
||||
assert list(Path(target.session_dir).glob("*.json")) == []
|
||||
|
||||
|
||||
def test_import_sqlite_write_failure_is_reported_and_preserves_existing_row(stores):
|
||||
source = stores("source", DefaultSessionStore)
|
||||
assert source.add_message("reference", "user", "New history")
|
||||
target = stores("target")
|
||||
assert target.add_message("reference", "user", "Original history")
|
||||
target._connect().execute("PRAGMA query_only=ON")
|
||||
report = target.import_sessions(source.export_all(), overwrite=True)
|
||||
assert report.imported == 0
|
||||
assert report.skipped_count == 1
|
||||
assert "write failed" in report.skipped[0]["reason"]
|
||||
target.invalidate_cache()
|
||||
assert target.get_chat_history("reference")[0]["content"] == "Original history"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reset_live_fields", [True, False])
|
||||
def test_restore_preserves_archived_history_and_tool_calls(stores, reset_live_fields):
|
||||
tool_calls = [{"id": "call-1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}]
|
||||
payload = {"sessions": [{
|
||||
"session_id": "reference",
|
||||
"gateway_session_id": "live",
|
||||
"agent_id": "agent",
|
||||
"metadata": {"gateway_session_id": "live", "agent_id": "agent", "lineage_id": "shared"},
|
||||
"archived_messages": [{"role": "user", "content": "Archived", "timestamp": 1}],
|
||||
"messages": [{"role": "assistant", "content": "Tool call", "tool_calls": tool_calls}],
|
||||
}]}
|
||||
target = stores("target", active_window=1, retention="truncate")
|
||||
assert target.add_message("existing", "user", "Existing")
|
||||
assert target.import_sessions(payload, reset_live_fields=reset_live_fields).imported == 1
|
||||
record = stores("target").export_session("reference")["sessions"][0]
|
||||
assert record["messages"][0]["tool_calls"] == tool_calls
|
||||
assert record["archived_messages"][0]["content"] == "Archived"
|
||||
assert record["metadata"]["lineage_id"] == "shared"
|
||||
if reset_live_fields:
|
||||
assert record.get("gateway_session_id") is None
|
||||
assert record.get("agent_id") is None
|
||||
assert "gateway_session_id" not in record["metadata"]
|
||||
assert "agent_id" not in record["metadata"]
|
||||
else:
|
||||
assert record["gateway_session_id"] == "live"
|
||||
assert record["agent_id"] == "agent"
|
||||
assert record["metadata"]["gateway_session_id"] == "live"
|
||||
assert record["metadata"]["agent_id"] == "agent"
|
||||
|
||||
|
||||
def test_sqlite_restore_keeps_import_guards(stores):
|
||||
source = stores("source", DefaultSessionStore)
|
||||
assert source.add_message("reference", "user", "Original")
|
||||
payload = source.export_all()
|
||||
payload["sessions"].append(dict(payload["sessions"][0]))
|
||||
target = stores("target")
|
||||
report = target.import_sessions(payload)
|
||||
assert report.imported == 1
|
||||
assert "duplicate" in report.skipped[0]["reason"]
|
||||
assert target.import_sessions(payload).imported == 0
|
||||
assert target.import_sessions(payload, overwrite=True).imported == 1
|
||||
assert target.export_session("missing")["sessions"] == []
|
||||
payload["version"] = target.PORTABLE_VERSION + 1
|
||||
assert target.import_sessions(payload, overwrite=True).imported == 0
|
||||
assert target.get_chat_history("reference")[0]["content"] == "Original"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("overwrite", [False, True])
|
||||
def test_import_does_not_clobber_peer_created_after_initial_check(stores, monkeypatch, overwrite):
|
||||
target = stores("target")
|
||||
peer = stores("target")
|
||||
original = target._save_imported_session
|
||||
|
||||
def create_peer_before_write(*args, **kwargs):
|
||||
assert peer.add_message("reference", "user", "peer-created")
|
||||
return original(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(target, "_save_imported_session", create_peer_before_write)
|
||||
payload = {"sessions": [{"session_id": "reference", "messages": [{"role": "user", "content": "imported"}]}]}
|
||||
report = target.import_sessions(payload, overwrite=overwrite)
|
||||
assert report.imported == int(overwrite)
|
||||
assert target.get_chat_history("reference")[0]["content"] == ("imported" if overwrite else "peer-created")
|
||||
if overwrite:
|
||||
assert target._cache["reference"].messages[0].content == "imported"
|
||||
else:
|
||||
assert report.skipped[0]["reason"] == "already exists (use overwrite)"
|
||||
|
||||
|
||||
def test_lineage_export_does_not_materialize_unrelated_transcripts(stores):
|
||||
store = stores("source")
|
||||
for sid in ("root", "continuation"):
|
||||
assert store.add_message(sid, "user", sid)
|
||||
assert store.update_session_metadata(sid, lineage_id="shared")
|
||||
for index in range(24):
|
||||
sid = f"unrelated-{index}"
|
||||
assert store.add_message(sid, "user", "x" * 256_000)
|
||||
assert store.update_session_metadata(sid, lineage_id="other", value=float("nan"))
|
||||
tracemalloc.start()
|
||||
try:
|
||||
payload = store.export_session("continuation")
|
||||
_, peak = tracemalloc.get_traced_memory()
|
||||
finally:
|
||||
tracemalloc.stop()
|
||||
assert {record["session_id"] for record in payload["sessions"]} == {"root", "continuation"}
|
||||
# The unrelated transcripts total over 6 MB. A single-record parse can
|
||||
# allocate its input and decoded content, but must not retain all rows.
|
||||
assert peak < 2_000_000
|
||||
Reference in new issue
Block a user