citadeldb-ms-agent-framework 2.0.0__tar.gz

citadeldb_ms_agent_framework-2.0.0/.gitignore

@@ -0,0 +1,18 @@
+/target
+**/*.rs.bk
+*.swp
+*.swo
+*~
+.DS_Store
+site/public/
+site/static/wasm/*.wasm
+site/static/wasm/*.js
+/notes/
+__pycache__/
+*.py[cod]
+# maturin build output; the .pyd is caught by the line above only by accident.
+*.pdb
+.pytest_cache/
+.mypy_cache/
+/dist/
+packaging/*/dist/

citadeldb_ms_agent_framework-2.0.0/PKG-INFO

@@ -0,0 +1,100 @@
+Metadata-Version: 2.5
+Name: citadeldb-ms-agent-framework
+Version: 2.0.0
+Summary: Microsoft Agent Framework chat history backed by Citadel: encrypted at rest, with deletes that destroy the key
+Project-URL: Homepage, https://citadeldb.dev
+Project-URL: Repository, https://github.com/yp3y5akh0v/citadel
+Author: Yuriy Peysakhov
+License-Expression: Apache-2.0
+Keywords: agent-framework,agents,autogen,chat-history,encryption,memory,microsoft-agent-framework
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Database
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: agent-framework-core<2,>=1.13
+Requires-Dist: citadeldb<3,>=2.0
+Provides-Extra: test
+Requires-Dist: pytest-asyncio>=0.23; extra == 'test'
+Requires-Dist: pytest>=8; extra == 'test'
+Description-Content-Type: text/markdown
+
+# citadeldb-ms-agent-framework
+
+[Microsoft Agent Framework](https://github.com/microsoft/agent-framework) storage backed
+by [Citadel](https://citadeldb.dev). Encrypted at rest, embedded in your process, and
+deletes that destroy the key, not just the row.
+
+```
+pip install citadeldb-ms-agent-framework
+```
+
+Two providers for two jobs, matching how the framework's own Redis integration is split:
+
+| Class | Implements | Use when |
+|---|---|---|
+| `CitadelHistoryProvider` | `HistoryProvider` | a session must recover its complete transcript |
+| `CitadelContextProvider` | `ContextProvider` | an agent should recall relevant facts across sessions |
+
+```python
+from agent_framework import Agent
+from citadeldb_ms_agent_framework import CitadelContextProvider, CitadelHistoryProvider
+
+agent = Agent(
+    client=chat_client,          # any agent_framework chat client
+    context_providers=[
+        CitadelHistoryProvider("agent.cdl", key="your-passphrase"),
+        CitadelContextProvider("agent.cdl", key="your-passphrase", scope="user-123"),
+    ],
+)
+```
+
+Both can share one encrypted file: a path already open on this thread, under the same
+passphrase, is shared. Construct them on the same thread.
+
+These are the framework's own extension points, with the file encrypted and a key per
+message. The built-in `FileHistoryProvider` writes plaintext JSONL or MessagePack.
+
+## Deletes destroy the key
+
+```python
+history = CitadelHistoryProvider("agent.cdl", key="your-passphrase")
+memory = CitadelContextProvider("agent.cdl", key="your-passphrase", scope="user-123")
+
+await history.forget("session-42")   # returns the number erased
+await memory.forget()                # this provider's whole scope
+```
+
+Clearing a conversation destroys each message's own key and then deletes its row, so any
+ciphertext surviving elsewhere stays unreadable.
+
+## History provider
+
+Implements `get_messages` and `save_messages`; the base class's `before_run`/`after_run`
+handle loading and storing according to its configuration flags, so an audit-only or
+evaluation-only provider works as documented:
+
+```python
+CitadelHistoryProvider("agent.cdl", key="your-passphrase", load_messages=False)  # stores, never loads
+```
+
+Messages round-trip through the framework's own serialization, so roles, author names,
+multi-part contents and `additional_properties` all survive.
+
+## Context provider
+
+Recalls with Citadel's hybrid search: vector distance, keyword rank and recency, fused
+into one score. The default `MockEmbedder` is lexical; pass `embedder=` a real one to
+match across wording.
+
+```python
+memory = CitadelContextProvider("agent.cdl", key="your-passphrase", scope="user-123", limit=5)
+```
+
+Memories are scoped rather than session-bound, so a later conversation can recall an
+earlier one. `scope` is the boundary an erasure request applies to.
+
+## License
+
+Apache-2.0

citadeldb_ms_agent_framework-2.0.0/README.md

@@ -0,0 +1,78 @@
+# citadeldb-ms-agent-framework
+
+[Microsoft Agent Framework](https://github.com/microsoft/agent-framework) storage backed
+by [Citadel](https://citadeldb.dev). Encrypted at rest, embedded in your process, and
+deletes that destroy the key, not just the row.
+
+```
+pip install citadeldb-ms-agent-framework
+```
+
+Two providers for two jobs, matching how the framework's own Redis integration is split:
+
+| Class | Implements | Use when |
+|---|---|---|
+| `CitadelHistoryProvider` | `HistoryProvider` | a session must recover its complete transcript |
+| `CitadelContextProvider` | `ContextProvider` | an agent should recall relevant facts across sessions |
+
+```python
+from agent_framework import Agent
+from citadeldb_ms_agent_framework import CitadelContextProvider, CitadelHistoryProvider
+
+agent = Agent(
+    client=chat_client,          # any agent_framework chat client
+    context_providers=[
+        CitadelHistoryProvider("agent.cdl", key="your-passphrase"),
+        CitadelContextProvider("agent.cdl", key="your-passphrase", scope="user-123"),
+    ],
+)
+```
+
+Both can share one encrypted file: a path already open on this thread, under the same
+passphrase, is shared. Construct them on the same thread.
+
+These are the framework's own extension points, with the file encrypted and a key per
+message. The built-in `FileHistoryProvider` writes plaintext JSONL or MessagePack.
+
+## Deletes destroy the key
+
+```python
+history = CitadelHistoryProvider("agent.cdl", key="your-passphrase")
+memory = CitadelContextProvider("agent.cdl", key="your-passphrase", scope="user-123")
+
+await history.forget("session-42")   # returns the number erased
+await memory.forget()                # this provider's whole scope
+```
+
+Clearing a conversation destroys each message's own key and then deletes its row, so any
+ciphertext surviving elsewhere stays unreadable.
+
+## History provider
+
+Implements `get_messages` and `save_messages`; the base class's `before_run`/`after_run`
+handle loading and storing according to its configuration flags, so an audit-only or
+evaluation-only provider works as documented:
+
+```python
+CitadelHistoryProvider("agent.cdl", key="your-passphrase", load_messages=False)  # stores, never loads
+```
+
+Messages round-trip through the framework's own serialization, so roles, author names,
+multi-part contents and `additional_properties` all survive.
+
+## Context provider
+
+Recalls with Citadel's hybrid search: vector distance, keyword rank and recency, fused
+into one score. The default `MockEmbedder` is lexical; pass `embedder=` a real one to
+match across wording.
+
+```python
+memory = CitadelContextProvider("agent.cdl", key="your-passphrase", scope="user-123", limit=5)
+```
+
+Memories are scoped rather than session-bound, so a later conversation can recall an
+earlier one. `scope` is the boundary an erasure request applies to.
+
+## License
+
+Apache-2.0

citadeldb_ms_agent_framework-2.0.0/pyproject.toml

@@ -0,0 +1,45 @@
+[build-system]
+requires = ["hatchling", "hatch-vcs"]
+build-backend = "hatchling.build"
+
+[project]
+# Named for the framework's owner: "agent framework" alone reads as a category.
+name = "citadeldb-ms-agent-framework"
+dynamic = ["version"]
+description = "Microsoft Agent Framework chat history backed by Citadel: encrypted at rest, with deletes that destroy the key"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "Apache-2.0"
+authors = [{ name = "Yuriy Peysakhov" }]
+keywords = [
+  "microsoft-agent-framework", "agent-framework", "autogen", "agents", "memory",
+  "chat-history", "encryption",
+]
+classifiers = [
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Developers",
+  "Programming Language :: Python :: 3",
+  "Topic :: Database",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+# A filtered recall answers `k` matching rows from 2.0; before it, a scope's
+# own memories could rank outside the scan and recall would return short.
+dependencies = ["citadeldb>=2.0,<3", "agent-framework-core>=1.13,<2"]
+
+[project.optional-dependencies]
+test = ["pytest>=8", "pytest-asyncio>=0.23"]
+
+[project.urls]
+Homepage = "https://citadeldb.dev"
+Repository = "https://github.com/yp3y5akh0v/citadel"
+
+# The version comes from the release tag, so there is nothing to bump.
+[tool.hatch.version]
+source = "vcs"
+raw-options = { root = "../..", tag_regex = '^v(?P<version>\d+\.\d+\.\d+)$' }
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/citadeldb_ms_agent_framework"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"

citadeldb_ms_agent_framework-2.0.0/src/citadeldb_ms_agent_framework/__init__.py

@@ -0,0 +1,14 @@
+"""Microsoft Agent Framework storage backed by Citadel, encrypted at rest."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from .history import CitadelHistoryProvider
+from .memory import CitadelContextProvider
+
+__all__ = ["CitadelContextProvider", "CitadelHistoryProvider", "__version__"]
+
+
+try:
+    __version__ = version("citadeldb-ms-agent-framework")
+except PackageNotFoundError:  # running from a source tree, never installed
+    __version__ = "0+unknown"

citadeldb_ms_agent_framework-2.0.0/src/citadeldb_ms_agent_framework/history.py

@@ -0,0 +1,164 @@
+"""HistoryProvider over an encrypted Citadel region."""
+from __future__ import annotations
+
+import asyncio
+from typing import Any, ClassVar, Sequence
+
+import citadeldb
+from agent_framework import HistoryProvider, Message
+
+KIND = "message"
+DEFAULT_PATH = "agent_history.cdl"
+DEFAULT_REGION = "history"
+PAGE = 10_000
+
+
+def _page(mem: Any, region: str, session_id: str) -> list[Any]:
+    """Page to the end: one fetch is bounded, and a partial erase must not look whole."""
+    out: list[Any] = []
+    after = None
+    while True:
+        got = mem.fetch(region, KIND, payload_filter={"sid": session_id}, limit=PAGE,
+                        after_id=after)
+        out.extend(got)
+        if len(got) < PAGE:
+            return out
+        after = got[-1].id
+# session_id is optional in the protocol; keep unattributed history together.
+DEFAULT_SESSION = "default"
+
+
+# A Database is pinned to its opening thread, so workers take Memory, not self.
+
+
+def _text_of(message: Message) -> str:
+    """An atom needs text to embed; a textless message falls back to role."""
+    return message.text or str(message.role)
+
+
+def _load(mem: Any, region: str, session_id: str) -> list[Message]:
+    hits = _page(mem, region, session_id)
+    return [Message.from_dict(h.payload["msg"]) for h in hits]
+
+
+def _append(mem: Any, region: str, session_id: str, messages: Sequence[Message]) -> None:
+    atoms = [
+        {
+            "kind": KIND,
+            "text": _text_of(m),
+            "payload": {"sid": session_id, "msg": m.to_dict()},
+        }
+        for m in messages
+    ]
+    if atoms:
+        # One batch draws a contiguous id range, so list order survives.
+        mem.remember_batch(region, atoms)
+
+
+def _forget(mem: Any, region: str, session_id: str) -> int:
+    hits = _page(mem, region, session_id)
+    if not hits:
+        return 0
+    return mem.forget(region, [h.id for h in hits]).erased_count
+
+
+def _recall(
+    mem: Any, region: str, session_id: str, query: str, limit: int
+) -> list[Message]:
+    hits = mem.recall(
+        region,
+        text=query,
+        k=limit,
+        kinds=[KIND],
+        options=citadeldb.RecallOptions(payload_filter={"sid": session_id}),
+    )
+    return [Message.from_dict(h.payload["msg"]) for h in hits]
+
+
+class CitadelHistoryProvider(HistoryProvider):
+    """A `HistoryProvider` backed by one encrypted Citadel region."""
+
+    DEFAULT_SOURCE_ID: ClassVar[str] = "citadel_history"
+
+    def __init__(
+        self,
+        path: str = DEFAULT_PATH,
+        key: str = "",
+        *,
+        source_id: str = DEFAULT_SOURCE_ID,
+        region: str = DEFAULT_REGION,
+        embedder: Any | None = None,
+        load_messages: bool = True,
+        store_inputs: bool = True,
+        store_context_messages: bool = False,
+        store_context_from: set[str] | None = None,
+        store_outputs: bool = True,
+    ) -> None:
+        super().__init__(
+            source_id=source_id,
+            load_messages=load_messages,
+            store_inputs=store_inputs,
+            store_context_messages=store_context_messages,
+            store_context_from=store_context_from,
+            store_outputs=store_outputs,
+        )
+        if not key:
+            raise ValueError("a passphrase is required: transcripts are the payload")
+        try:
+            self._db = citadeldb.connect(path, key=key, region_keys=True)
+        except citadeldb.OperationalError as e:
+            if "locked" not in str(e):
+                raise
+            raise RuntimeError(
+                f"{path} is open in another process. Citadel is embedded, so one "
+                f"process owns the file."
+            ) from e
+        self._mem = self._db.memory()
+        self._region = region
+        # Idempotent for a region of the same width, so a dim clash raises here.
+        self._mem.create_encrypted_region(
+            region, embedder or citadeldb.MockEmbedder(dim=64)
+        )
+
+    # ---- the abstract surface --------------------------------------------
+    # The bindings are sync, so a worker thread keeps the event loop free.
+
+    async def get_messages(
+        self, session_id: str | None, *, state: dict[str, Any] | None = None, **kwargs: Any
+    ) -> list[Message]:
+        return await asyncio.to_thread(
+            _load, self._mem, self._region, session_id or DEFAULT_SESSION
+        )
+
+    async def save_messages(
+        self,
+        session_id: str | None,
+        messages: Sequence[Message],
+        *,
+        state: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        # Appends rather than replaces: history is a transcript, not a set.
+        await asyncio.to_thread(
+            _append,
+            self._mem,
+            self._region,
+            session_id or DEFAULT_SESSION,
+            list(messages),
+        )
+
+    # ---- beyond the protocol ---------------------------------------------
+
+    async def search(
+        self, session_id: str | None, query: str, *, limit: int = 5
+    ) -> list[Message]:
+        """Messages from one session ranked by hybrid recall, best first."""
+        return await asyncio.to_thread(
+            _recall, self._mem, self._region, session_id or DEFAULT_SESSION, query, limit
+        )
+
+    async def forget(self, session_id: str | None) -> int:
+        """Destroy one session's messages, returning the number erased."""
+        return await asyncio.to_thread(
+            _forget, self._mem, self._region, session_id or DEFAULT_SESSION
+        )

citadeldb_ms_agent_framework-2.0.0/src/citadeldb_ms_agent_framework/memory.py

@@ -0,0 +1,163 @@
+"""ContextProvider over an encrypted Citadel region."""
+from __future__ import annotations
+
+import asyncio
+from typing import Any, ClassVar, Sequence
+
+import citadeldb
+from agent_framework import ContextProvider, Message
+
+KIND = "memory"
+DEFAULT_PATH = "agent_memory.cdl"
+DEFAULT_REGION = "memories"
+PAGE = 10_000
+
+
+def _page(mem: Any, region: str, criterion: dict[str, Any]) -> list[Any]:
+    """Page to the end: one fetch is bounded, and a partial erase must not look whole."""
+    out: list[Any] = []
+    after = None
+    while True:
+        got = mem.fetch(region, KIND, payload_filter=criterion, limit=PAGE, after_id=after)
+        out.extend(got)
+        if len(got) < PAGE:
+            return out
+        after = got[-1].id
+# Roles worth remembering; tool traffic is transcript detail, not knowledge.
+_REMEMBERED_ROLES = ("user", "assistant", "system")
+
+
+# A Database is pinned to its opening thread, so workers take Memory, not self.
+
+
+def _role_of(message: Message) -> str:
+    role = message.role
+    return getattr(role, "value", None) or str(role)
+
+
+def _remember(mem: Any, region: str, scope: str, texts: Sequence[str]) -> None:
+    atoms = [
+        {"kind": KIND, "text": t, "payload": {"scope": scope, "text": t}} for t in texts
+    ]
+    if atoms:
+        mem.remember_batch(region, atoms)
+
+
+def _recall(mem: Any, region: str, scope: str, query: str, limit: int) -> list[str]:
+    """Distinct memories for `scope`, best first.
+
+    after_run stores every turn verbatim, so a fact the user restates is stored
+    once per turn. Those copies are one memory to the model, and asking the
+    engine for `limit` rows would spend the whole budget on them - 30 repeats of
+    one fact deliver one line and hide everything else in the scope. Widen until
+    `limit` distinct texts are found or the scope runs out.
+    """
+    options = citadeldb.RecallOptions(payload_filter={"scope": scope})
+    k = max(limit, 32)
+    while True:
+        hits = mem.recall(region, text=query, k=k, kinds=[KIND], options=options)
+        out: list[str] = []
+        seen: set[str] = set()
+        for h in hits:
+            text = h.payload["text"]
+            if text in seen:
+                continue
+            seen.add(text)
+            out.append(text)
+        if len(out) >= limit or len(hits) < k:
+            return out[:limit]
+        k *= 2
+
+
+def _forget(mem: Any, region: str, scope: str) -> int:
+    hits = _page(mem, region, {"scope": scope})
+    if not hits:
+        return 0
+    return mem.forget(region, [h.id for h in hits]).erased_count
+
+
+class CitadelContextProvider(ContextProvider):
+    """A `ContextProvider` backed by one encrypted Citadel region."""
+
+    DEFAULT_SOURCE_ID: ClassVar[str] = "citadel_memory"
+    DEFAULT_CONTEXT_PROMPT: ClassVar[str] = (
+        "## Memories\nConsider the following memories from earlier conversations:"
+    )
+
+    def __init__(
+        self,
+        path: str = DEFAULT_PATH,
+        key: str = "",
+        *,
+        source_id: str = DEFAULT_SOURCE_ID,
+        scope: str = "default",
+        region: str = DEFAULT_REGION,
+        embedder: Any | None = None,
+        limit: int = 5,
+        context_prompt: str = DEFAULT_CONTEXT_PROMPT,
+    ) -> None:
+        super().__init__(source_id)
+        if not key:
+            raise ValueError("a passphrase is required: memories are the payload")
+        self.scope = scope
+        self.limit = limit
+        self.context_prompt = context_prompt
+        try:
+            self._db = citadeldb.connect(path, key=key, region_keys=True)
+        except citadeldb.OperationalError as e:
+            if "locked" not in str(e):
+                raise
+            raise RuntimeError(
+                f"{path} is open in another process. Citadel is embedded, so one "
+                f"process owns the file."
+            ) from e
+        self._mem = self._db.memory()
+        self._region = region
+        # Idempotent for a region of the same width, so a dim clash raises here.
+        self._mem.create_encrypted_region(
+            region, embedder or citadeldb.MockEmbedder(dim=64)
+        )
+
+    # ---- the pipeline hooks ----------------------------------------------
+
+    async def before_run(
+        self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]
+    ) -> None:
+        """Recall what is relevant to this turn and add it to the context."""
+        query = "\n".join(
+            m.text for m in context.input_messages if m and m.text and m.text.strip()
+        )
+        if not query:
+            return
+        memories = await asyncio.to_thread(
+            _recall, self._mem, self._region, self.scope, query, self.limit
+        )
+        if not memories:
+            return
+        context.extend_messages(
+            self.source_id,
+            [Message("user", [f"{self.context_prompt}\n" + "\n".join(memories)])],
+        )
+
+    async def after_run(
+        self, *, agent: Any, session: Any, context: Any, state: dict[str, Any]
+    ) -> None:
+        """Remember this turn, inputs and response alike."""
+        turn: list[Message] = list(context.input_messages)
+        if context.response and context.response.messages:
+            turn.extend(context.response.messages)
+        texts = [
+            m.text
+            for m in turn
+            if m and m.text and m.text.strip() and _role_of(m) in _REMEMBERED_ROLES
+        ]
+        if texts:
+            await asyncio.to_thread(
+                _remember, self._mem, self._region, self.scope, texts
+            )
+
+    # ---- beyond the pipeline ---------------------------------------------
+
+    async def forget(self) -> int:
+        """Destroy this scope's memories, returning the number erased."""
+        return await asyncio.to_thread(_forget, self._mem, self._region, self.scope)

citadeldb_ms_agent_framework-2.0.0/tests/test_history.py

@@ -0,0 +1,366 @@
+import pytest
+from agent_framework import (
+    AgentSession,
+    ContextProvider,
+    FileHistoryProvider,
+    HistoryProvider,
+    InMemoryHistoryProvider,
+    Message,
+    SessionContext,
+)
+
+from citadeldb_ms_agent_framework import CitadelHistoryProvider
+
+
+@pytest.fixture(scope="module")
+def path(tmp_path_factory):
+    # Citadel takes an exclusive lock, so the whole module shares one file.
+    return str(tmp_path_factory.mktemp("maf") / "h.cdl")
+
+
+@pytest.fixture()
+def provider(path, request):
+    return CitadelHistoryProvider(path, "pw", source_id=request.node.name)
+
+
+def sid(request) -> str:
+    return request.node.name
+
+
+def msg(text, role="user", **kw):
+    return Message(role, [text], **kw)
+
+
+def context(session_id, inputs):
+    return SessionContext(session_id=session_id, input_messages=list(inputs))
+
+
+# ---- conformance ---------------------------------------------------------
+
+
+def test_is_a_history_provider(provider):
+    assert isinstance(provider, HistoryProvider)
+    assert isinstance(provider, ContextProvider)
+
+
+def test_configuration_flags_reach_the_base_class(path):
+    """before_run/after_run are the base class's; they read these flags."""
+    p = CitadelHistoryProvider(
+        path, "pw", source_id="flags", load_messages=False, store_outputs=False,
+        store_context_messages=True, store_context_from={"other"},
+    )
+    assert p.load_messages is False
+    assert p.store_outputs is False
+    assert p.store_context_messages is True
+    assert p.store_context_from == {"other"}
+    assert p.store_inputs is True
+
+
+def test_source_id_defaults_to_a_named_constant(path):
+    p = CitadelHistoryProvider(path, "pw")
+    assert p.source_id == CitadelHistoryProvider.DEFAULT_SOURCE_ID
+
+
+# ---- round-trip ----------------------------------------------------------
+
+
+async def test_messages_round_trip_in_order(provider, request):
+    s = sid(request)
+    await provider.save_messages(s, [msg("one"), msg("two", "assistant")])
+    await provider.save_messages(s, [msg("three")])
+    got = await provider.get_messages(s)
+    assert [m.text for m in got] == ["one", "two", "three"]
+
+
+async def test_roles_and_author_survive(provider, request):
+    s = sid(request)
+    await provider.save_messages(s, [
+        msg("sys", "system"),
+        msg("q", "user", author_name="alice"),
+        msg("a", "assistant"),
+    ])
+    got = await provider.get_messages(s)
+    assert [str(m.role) for m in got] == [str(m.role) for m in got]  # stable
+    assert got[1].author_name == "alice"
+
+
+async def test_additional_properties_survive(provider, request):
+    s = sid(request)
+    await provider.save_messages(
+        s, [Message("user", ["x"], additional_properties={"_excluded": True})]
+    )
+    got = await provider.get_messages(s)
+    assert got[0].additional_properties["_excluded"] is True
+
+
+async def test_multi_content_messages_round_trip(provider, request):
+    s = sid(request)
+    await provider.save_messages(s, [Message("user", ["first", "second"])])
+    got = await provider.get_messages(s)
+    assert len(got[0].contents) == 2
+
+
+# ---- semantics against the built-in providers ----------------------------
+
+
+async def test_save_appends_rather_than_replaces(provider, request):
+    """Built-in providers extend the stored list; history is a transcript."""
+    s = sid(request)
+    reference = InMemoryHistoryProvider()
+    state: dict = {}
+    for batch in ([msg("a")], [msg("b")]):
+        await provider.save_messages(s, batch)
+        await reference.save_messages(s, batch, state=state)
+    ours = [m.text for m in await provider.get_messages(s)]
+    theirs = [m.text for m in await reference.get_messages(s, state=state)]
+    assert ours == theirs == ["a", "b"]
+
+
+async def test_saving_nothing_is_not_an_error(provider, request):
+    s = sid(request)
+    await provider.save_messages(s, [])
+    assert await provider.get_messages(s) == []
+
+
+async def test_an_unseen_session_is_empty_not_an_error(provider):
+    assert await provider.get_messages("never-used") == []
+
+
+# ---- before_run and after_run --------------------------------------------
+
+
+async def test_before_run_loads_stored_history_into_context(provider, request):
+    s = sid(request)
+    await provider.save_messages(s, [msg("remembered")])
+    ctx = context(s, [msg("new question")])
+    await provider.before_run(
+        agent=None, session=AgentSession(session_id=s), context=ctx, state={}
+    )
+    loaded = ctx.get_messages(sources={provider.source_id})
+    assert [m.text for m in loaded] == ["remembered"]
+
+
+async def test_after_run_stores_the_input_messages(provider, request):
+    """store_inputs defaults True, so a turn's inputs land in the transcript."""
+    s = sid(request)
+    ctx = context(s, [msg("asked")])
+    await provider.after_run(
+        agent=None, session=AgentSession(session_id=s), context=ctx, state={}
+    )
+    assert [m.text for m in await provider.get_messages(s)] == ["asked"]
+
+
+async def test_a_full_turn_accumulates(provider, request):
+    s = sid(request)
+    session = AgentSession(session_id=s)
+    for text in ("first", "second"):
+        ctx = context(s, [msg(text)])
+        await provider.before_run(agent=None, session=session, context=ctx, state={})
+        await provider.after_run(agent=None, session=session, context=ctx, state={})
+    assert [m.text for m in await provider.get_messages(s)] == ["first", "second"]
+
+
+async def test_load_messages_false_still_stores(path):
+    """An audit-only provider: stores but never loads."""
+    p = CitadelHistoryProvider(path, "pw", source_id="audit", load_messages=False)
+    ctx = context("audit-s", [msg("recorded")])
+    await p.after_run(
+        agent=None, session=AgentSession(session_id="audit-s"), context=ctx, state={}
+    )
+    assert [m.text for m in await p.get_messages("audit-s")] == ["recorded"]
+
+
+async def test_store_inputs_false_records_nothing_from_inputs(path):
+    p = CitadelHistoryProvider(path, "pw", source_id="nostore", store_inputs=False)
+    ctx = context("nostore-s", [msg("ignored")])
+    await p.after_run(
+        agent=None, session=AgentSession(session_id="nostore-s"), context=ctx, state={}
+    )
+    assert await p.get_messages("nostore-s") == []
+
+
+@pytest.mark.filterwarnings("ignore:.*FileHistoryProvider is experimental.*")
+async def test_matches_the_file_provider_on_a_turn(provider, request, tmp_path):
+    """FileHistoryProvider is the closest built-in analogue; behave like it."""
+    s = sid(request)
+    reference = FileHistoryProvider(str(tmp_path / "hist"), source_id="ref")
+    for p in (provider, reference):
+        ctx = context(s, [msg("a"), msg("b", "assistant")])
+        await p.after_run(
+            agent=None, session=AgentSession(session_id=s), context=ctx, state={}
+        )
+    ours = [m.text for m in await provider.get_messages(s)]
+    theirs = [m.text for m in await reference.get_messages(s)]
+    assert ours == theirs
+
+
+async def test_state_is_ignored_like_the_file_provider(provider, request):
+    """State-backed storage is InMemoryHistoryProvider's model, not ours."""
+    s = sid(request)
+    await provider.save_messages(s, [msg("stored")], state={"messages": []})
+    got = await provider.get_messages(s, state={"messages": [msg("phantom")]})
+    assert [m.text for m in got] == ["stored"]
+
+
+# ---- session scoping -----------------------------------------------------
+
+
+async def test_search_finds_a_session_buried_under_another(provider):
+    """Discarding after the scan spends the budget on the busy session."""
+    quiet, noisy = "buried-quiet", "buried-noisy"
+    # Ranked above the target for this query, and enough of them to fill the scan.
+    await provider.save_messages(
+        noisy, [msg(f"why did the release break run {i}") for i in range(60)]
+    )
+    await provider.save_messages(
+        quiet, [msg("the deployment failed because the disk was full")]
+    )
+    hits = await provider.search(quiet, "why did the release break?", limit=1)
+    assert hits, "the quiet session's only match was crowded out"
+    assert "disk was full" in hits[0].text
+
+
+async def test_sessions_do_not_see_each_other(provider):
+    await provider.save_messages("iso-a", [msg("mine")])
+    await provider.save_messages("iso-b", [msg("yours")])
+    assert [m.text for m in await provider.get_messages("iso-a")] == ["mine"]
+
+
+async def test_session_ids_are_matched_exactly_not_by_prefix(provider):
+    await provider.save_messages("pre", [msg("outer")])
+    await provider.save_messages("pre-fix", [msg("inner")])
+    assert [m.text for m in await provider.get_messages("pre")] == ["outer"]
+
+
+async def test_a_none_session_id_uses_one_stable_bucket(provider):
+    """The protocol allows None; the file provider uses a fixed stem too."""
+    await provider.save_messages(None, [msg("unattributed")])
+    assert [m.text for m in await provider.get_messages(None)] == ["unattributed"]
+
+
+# ---- erasure -------------------------------------------------------------
+
+
+async def test_forget_destroys_one_session(provider):
+    await provider.save_messages("gone", [msg("a"), msg("b")])
+    await provider.save_messages("kept", [msg("c")])
+    assert await provider.forget("gone") == 2
+    assert await provider.get_messages("gone") == []
+    assert len(await provider.get_messages("kept")) == 1
+
+
+async def test_forgetting_nothing_is_zero_not_an_error(provider):
+    assert await provider.forget("never-existed") == 0
+
+
+async def test_a_session_is_reusable_after_forgetting(provider):
+    await provider.save_messages("reuse", [msg("first")])
+    await provider.forget("reuse")
+    await provider.save_messages("reuse", [msg("second")])
+    assert [m.text for m in await provider.get_messages("reuse")] == ["second"]
+
+
+async def test_it_survives_a_reopen(tmp_path):
+    """A provider that cannot be reopened is not persistence."""
+    import gc
+
+    p = str(tmp_path / "reopen.cdl")
+    first = CitadelHistoryProvider(p, "pw")
+    await first.save_messages("s", [msg("the disk was full"), msg("ok", "assistant")])
+    del first
+    gc.collect()
+
+    again = CitadelHistoryProvider(p, "pw")
+    assert [m.text for m in await again.get_messages("s")] == [
+        "the disk was full", "ok"
+    ]
+
+
+async def test_concurrent_saves_all_land(tmp_path):
+    """Many in-flight sessions means many worker threads against one engine."""
+    import asyncio
+
+    h = CitadelHistoryProvider(str(tmp_path / "conc.cdl"), "pw")
+    await asyncio.gather(*(h.save_messages(f"s{i}", [msg(f"m{i}")]) for i in range(40)))
+    total = 0
+    for i in range(40):
+        total += len(await h.get_messages(f"s{i}"))
+    assert total == 40
+
+
+async def test_a_wrong_passphrase_cannot_reopen(tmp_path):
+    """Transcripts are the payload, so the encryption claim is pinned here."""
+    import gc
+
+    import citadeldb
+
+    p = str(tmp_path / "enc.cdl")
+    first = CitadelHistoryProvider(p, "right")
+    await first.save_messages("s", [msg("secret")])
+    del first
+    gc.collect()
+
+    with pytest.raises(citadeldb.EncryptionError):
+        CitadelHistoryProvider(p, "wrong")
+
+
+def test_a_passphrase_is_required(tmp_path):
+    with pytest.raises(ValueError, match="passphrase"):
+        CitadelHistoryProvider(str(tmp_path / "k.cdl"), "")
+
+
+# ---- edges ---------------------------------------------------------------
+
+
+async def test_a_message_with_no_text_is_storable(provider, request):
+    """An atom needs text to embed, so a textless message needs a handle."""
+    s = sid(request)
+    await provider.save_messages(s, [Message("assistant", [])])
+    got = await provider.get_messages(s)
+    assert len(got) == 1 and got[0].text == ""
+
+
+async def test_a_long_history_keeps_its_order(provider, request):
+    s = sid(request)
+    await provider.save_messages(s, [msg(f"turn {i}") for i in range(200)])
+    assert [m.text for m in await provider.get_messages(s)] == [
+        f"turn {i}" for i in range(200)
+    ]
+
+
+async def test_order_survives_across_separate_batches(provider, request):
+    s = sid(request)
+    for i in range(20):
+        await provider.save_messages(s, [msg(f"m{i}")])
+    assert [m.text for m in await provider.get_messages(s)] == [
+        f"m{i}" for i in range(20)
+    ]
+
+
+async def test_two_providers_share_one_database_file(tmp_path):
+    p = str(tmp_path / "shared.cdl")
+    a = CitadelHistoryProvider(p, "pw", source_id="a")
+    b = CitadelHistoryProvider(p, "pw", source_id="b")
+    await a.save_messages("s", [msg("written by a")])
+    assert [m.text for m in await b.get_messages("s")] == ["written by a"]
+
+
+async def test_the_event_loop_is_not_blocked(provider, request):
+    """The abstract methods are async, so sync bindings run off the loop."""
+    import asyncio
+
+    s = sid(request)
+    ticks = 0
+
+    async def tick():
+        nonlocal ticks
+        while True:
+            ticks += 1
+            await asyncio.sleep(0)
+
+    ticker = asyncio.create_task(tick())
+    await asyncio.sleep(0)
+    await provider.save_messages(s, [msg(f"loop {i}") for i in range(40)])
+    await provider.get_messages(s)
+    ticker.cancel()
+    assert ticks > 1, "the loop made no progress during a provider call"

citadeldb_ms_agent_framework-2.0.0/tests/test_memory.py

@@ -0,0 +1,264 @@
+import pytest
+from agent_framework import AgentSession, ContextProvider, Message, SessionContext
+
+from citadeldb_ms_agent_framework import CitadelContextProvider
+
+
+@pytest.fixture(scope="module")
+def path(tmp_path_factory):
+    # Citadel takes an exclusive lock, so the whole module shares one file.
+    return str(tmp_path_factory.mktemp("mem") / "m.cdl")
+
+
+@pytest.fixture()
+def provider(path, request):
+    return CitadelContextProvider(
+        path, "pw", source_id=request.node.name, scope=request.node.name
+    )
+
+
+def msg(text, role="user"):
+    return Message(role, [text])
+
+
+def context(inputs, session_id="s"):
+    return SessionContext(session_id=session_id, input_messages=list(inputs))
+
+
+async def turn(provider, texts, session_id="s"):
+    ctx = context([msg(t) for t in texts], session_id)
+    await provider.after_run(
+        agent=None, session=AgentSession(session_id=session_id), context=ctx, state={}
+    )
+
+
+# ---- conformance ---------------------------------------------------------
+
+
+def test_is_a_context_provider(provider):
+    assert isinstance(provider, ContextProvider)
+    assert provider.source_id
+
+
+def test_a_passphrase_is_required(tmp_path):
+    with pytest.raises(ValueError, match="passphrase"):
+        CitadelContextProvider(str(tmp_path / "k.cdl"), "")
+
+
+# ---- recall --------------------------------------------------------------
+
+
+async def test_a_restated_fact_does_not_crowd_out_the_rest_of_the_scope(provider):
+    """after_run stores every turn verbatim, so a fact the user restates is
+    stored once per turn. Those copies are one memory to the model, so asking
+    the engine for `limit` rows spends the whole budget on them."""
+    for _ in range(40):
+        await turn(provider, ["the deploy failed because the disk was full"])
+    for n in range(4):
+        await turn(provider, [f"unrelated note {n}"])
+
+    ctx = context([msg("why did the release break?")])
+    await provider.before_run(
+        agent=None, session=AgentSession(session_id="s"), context=ctx, state={}
+    )
+    delivered = "\n".join(m.text for m in ctx.get_messages(sources={provider.source_id}))
+    assert delivered.count("the disk was full") == 1, "duplicates reached the model"
+    for n in range(4):
+        assert f"unrelated note {n}" in delivered, (
+            f"note {n} was crowded out by 40 copies of one fact"
+        )
+
+
+async def test_a_memory_is_recalled_into_the_context(provider):
+    await turn(provider, ["the deploy failed because the disk was full"])
+    ctx = context([msg("why did the release break?")])
+    await provider.before_run(
+        agent=None, session=AgentSession(session_id="s"), context=ctx, state={}
+    )
+    added = ctx.get_messages(sources={provider.source_id})
+    assert added, "nothing was recalled"
+    assert "disk was full" in added[0].text
+
+
+async def test_a_scope_buried_under_another_is_still_recalled(path, request):
+    """Discarding after the scan spends the budget on the crowded scope."""
+    name = request.node.name
+    noisy = CitadelContextProvider(path, "pw", source_id=f"{name}-n", scope=f"{name}-n")
+    quiet = CitadelContextProvider(path, "pw", source_id=name, scope=name)
+    # Ranked above the target for this query, and enough of them to fill the scan.
+    await turn(noisy, [f"why did the release break run {i}" for i in range(60)])
+    await turn(quiet, ["the deploy failed because the disk was full"])
+
+    ctx = context([msg("why did the release break?")])
+    await quiet.before_run(
+        agent=None, session=AgentSession(session_id="s"), context=ctx, state={}
+    )
+    added = ctx.get_messages(sources={quiet.source_id})
+    assert added, "the quiet scope's only memory was crowded out"
+    assert "disk was full" in added[0].text
+
+
+async def test_the_context_prompt_frames_the_memories(provider):
+    await turn(provider, ["remembered fact"])
+    ctx = context([msg("recall")])
+    await provider.before_run(
+        agent=None, session=AgentSession(session_id="s"), context=ctx, state={}
+    )
+    added = ctx.get_messages(sources={provider.source_id})
+    assert added[0].text.startswith(CitadelContextProvider.DEFAULT_CONTEXT_PROMPT)
+
+
+async def test_recall_is_capped_by_limit(path):
+    p = CitadelContextProvider(path, "pw", source_id="cap", scope="cap", limit=2)
+    await turn(p, [f"fact number {i}" for i in range(6)])
+    ctx = context([msg("fact")])
+    await p.before_run(
+        agent=None, session=AgentSession(session_id="s"), context=ctx, state={}
+    )
+    added = ctx.get_messages(sources={p.source_id})
+    recalled = added[0].text.removeprefix(
+        CitadelContextProvider.DEFAULT_CONTEXT_PROMPT
+    ).strip().splitlines()
+    assert len(recalled) == 2, recalled
+
+
+async def test_a_repeated_fact_does_not_spend_the_whole_budget(path):
+    """A fact repeated across turns is stored each time, and it is one fact."""
+    p = CitadelContextProvider(path, "pw", source_id="dedup", scope="dedup", limit=5)
+    for _ in range(5):
+        await turn(p, ["my dog is called Mochi"])
+    await turn(p, ["I live in Berlin"])
+
+    ctx = context([msg("what do you know about me?")])
+    await p.before_run(
+        agent=None, session=AgentSession(session_id="s"), context=ctx, state={}
+    )
+    recalled = ctx.get_messages(sources={p.source_id})[0].text.removeprefix(
+        CitadelContextProvider.DEFAULT_CONTEXT_PROMPT
+    ).strip().splitlines()
+    assert len(recalled) == len(set(recalled)), recalled
+    assert any("Mochi" in r for r in recalled)
+    assert any("Berlin" in r for r in recalled)
+
+
+async def test_nothing_is_added_when_there_is_nothing_to_recall(provider):
+    ctx = context([msg("a question with no history behind it")])
+    await provider.before_run(
+        agent=None, session=AgentSession(session_id="s"), context=ctx, state={}
+    )
+    assert ctx.get_messages(sources={provider.source_id}) == []
+
+
+async def test_an_empty_input_recalls_nothing(provider):
+    await turn(provider, ["something"])
+    ctx = context([Message("user", [])])
+    await provider.before_run(
+        agent=None, session=AgentSession(session_id="s"), context=ctx, state={}
+    )
+    assert ctx.get_messages(sources={provider.source_id}) == []
+
+
+# ---- what gets remembered ------------------------------------------------
+
+
+async def test_memories_outlive_the_session_that_made_them(provider):
+    """Cross-session recall is this provider's job, not the history one's."""
+    await turn(provider, ["learned in session one"], session_id="one")
+    ctx = context([msg("learned")], session_id="two")
+    await provider.before_run(
+        agent=None, session=AgentSession(session_id="two"), context=ctx, state={}
+    )
+    added = ctx.get_messages(sources={provider.source_id})
+    assert added and "session one" in added[0].text
+
+
+async def test_scopes_do_not_see_each_other(path):
+    a = CitadelContextProvider(path, "pw", source_id="sa", scope="alice")
+    b = CitadelContextProvider(path, "pw", source_id="sb", scope="bob")
+    await turn(a, ["alice's private note"])
+    ctx = context([msg("private note")])
+    await b.before_run(
+        agent=None, session=AgentSession(session_id="s"), context=ctx, state={}
+    )
+    assert ctx.get_messages(sources={b.source_id}) == []
+
+
+async def test_a_turn_with_no_text_remembers_nothing(provider):
+    ctx = context([Message("user", [])])
+    await provider.after_run(
+        agent=None, session=AgentSession(session_id="s"), context=ctx, state={}
+    )
+    assert await provider.forget() == 0
+
+
+# ---- erasure -------------------------------------------------------------
+
+
+async def test_forget_destroys_the_scope(provider):
+    await turn(provider, ["a", "b"])
+    assert await provider.forget() == 2
+    ctx = context([msg("a")])
+    await provider.before_run(
+        agent=None, session=AgentSession(session_id="s"), context=ctx, state={}
+    )
+    assert ctx.get_messages(sources={provider.source_id}) == []
+
+
+async def test_it_survives_a_reopen(tmp_path):
+    """A region's embedder lives in memory, so recall must survive reopen."""
+    import gc
+
+    p = str(tmp_path / "reopen.cdl")
+    first = CitadelContextProvider(p, "pw", scope="u")
+    await turn(first, ["the deployment failed because the disk was full"])
+    del first
+    gc.collect()
+
+    again = CitadelContextProvider(p, "pw", scope="u")
+    ctx = context([msg("why did the release break?")])
+    await again.before_run(
+        agent=None, session=AgentSession(session_id="s"), context=ctx, state={}
+    )
+    assert ctx.get_messages(sources={again.source_id}), "recall did not survive reopen"
+
+
+async def test_forgetting_nothing_is_zero_not_an_error(provider):
+    assert await provider.forget() == 0
+
+
+# ---- the two providers together ------------------------------------------
+
+
+async def test_it_shares_a_database_with_the_history_provider(tmp_path):
+    """The pairing the framework's own Redis integration ships, on one file."""
+    from citadeldb_ms_agent_framework import CitadelHistoryProvider
+
+    p = str(tmp_path / "both.cdl")
+    memory = CitadelContextProvider(p, "pw", scope="u1")
+    history = CitadelHistoryProvider(p, "pw")
+    await turn(memory, ["remembered across sessions"])
+    await history.save_messages("s", [msg("this exact turn")])
+    assert len(await history.get_messages("s")) == 1
+    assert await memory.forget() == 1
+
+
+async def test_the_event_loop_is_not_blocked(provider):
+    import asyncio
+
+    ticks = 0
+
+    async def tick():
+        nonlocal ticks
+        while True:
+            ticks += 1
+            await asyncio.sleep(0)
+
+    ticker = asyncio.create_task(tick())
+    await asyncio.sleep(0)
+    await turn(provider, [f"fact {i}" for i in range(40)])
+    ctx = context([msg("fact")])
+    await provider.before_run(
+        agent=None, session=AgentSession(session_id="s"), context=ctx, state={}
+    )
+    ticker.cancel()
+    assert ticks > 1, "the loop made no progress during a provider call"