Add AgentMemory preparation

docs/guides/Agents.md

@@ -20,7 +20,10 @@ The agents API is intentionally built as a thin layer over existing Draive runti
 
 - `AgentIdentity` describes the agent instance: `uri`, `name`, `description`, `meta`.
 - `AgentMessage` is the fully prepared input payload: `thread`, `created`, `content`, `meta`.
-- `AgentThread` is the scoped runtime state propagated through `ctx.scope(...)`.
+- `AgentThread` is the scoped runtime state propagated through `ctx.scope(...)`: its `identifier`
+    names the conversation thread shared across nested agent calls, while `agent_uri` marks the
+    URI of the agent currently executing within it - agents bind a thread stamped with their own
+    URI while handling a message.
 - `AgentExecuting` is the executor protocol:
     `AgentMessage -> AsyncIterable[MultimodalContentPart | ProcessingEvent]`.
 
@@ -171,9 +174,13 @@ APIs, or provide the required context explicitly.
 
 ### Persist Context Across Turns With `AgentMemory`
 
-`Agent.generative(...)`, `Agent.steps(...)` and `Agent.from_skill(...)` accept a `memory` argument
-controlling how model context is recalled before each turn and persisted afterwards. The default is
-`AgentMemory.disabled`, which scopes context to a single turn.
+`Agent.generative(...)` and `Agent.from_skill(...)` accept a `memory` argument controlling how
+model context is prepared and recalled before each turn and persisted afterwards. The default is
+`AgentMemory.disabled`, which scopes context to a single turn. `Agent.steps(...)` applies no memory
+implicitly - compose the predefined memory steps (`memory.prepare_step(...)`,
+`memory.recall_step()`, `memory.remember_step()`) into the pipeline where needed;
+`Agent.generative(...)` and `Agent.from_skill(...)` invoke `memory.prepare`, `memory.recall`, and
+`memory.remember` directly within their step bodies instead.
 
 ```python
 from draive import Agent, AgentMemory
@@ -185,20 +192,32 @@ assistant: Agent = Agent.generative(
 )
 ```
 
-Memory is keyed by the active `AgentThread`, so one memory instance serves multiple concurrent
-conversation threads. Context is stored as the latest snapshot per thread: whatever is remembered
-after a turn becomes the next recalled context, exactly as provided, which allows steps to compact,
-summarize, or replace the context freely.
-
-- `AgentMemory.volatile(...)` keeps snapshots in-process, with optional LRU eviction via
-    `threads_limit`; intended for local development, tests, and single-process deployments.
-- `PostgresAgentMemory.prepare(identity)` (from `draive.postgres`) persists immutable snapshots in
-    PostgreSQL, keyed by agent identity and thread; run `PostgresAgentMemory.migrate()` once to
-    create its schema.
-- `AgentMemory(recalling=..., remembering=...)` wraps custom async callables for any other backend.
-
-Each memory instance is intended to serve exactly one agent: state is scoped per thread only, so
-sharing an instance between agents would mix their contexts within a thread.
+Memory operations receive the active `AgentThread` - the executing agent's URI together with the
+conversation thread identifier - so one memory instance can serve multiple agents and multiple
+concurrent conversation threads, as long as the implementation keys stored context by both (the
+built-in ones do). Context is stored as the latest snapshot per agent and thread: whatever is
+remembered after a turn becomes the next recalled context, exactly as provided, which allows steps
+to compact, summarize, or replace the context freely.
+
+- `AgentMemory.volatile(...)` keeps snapshots in-process, with optional LRU-like eviction via
+    `threads_limit`. An entry's usage is refreshed by `prepare` and `remember`, but not by
+    `recall`; intended for local development, tests, and single-process deployments.
+- `PostgresAgentMemory.instance()` (from `draive.postgres`) persists immutable snapshots in
+    PostgreSQL, keyed by executing agent URI and thread; run `PostgresAgentMemory.migrate()` once
+    to create its schema.
+- `AgentMemory(recalling=..., remembering=...)` wraps custom async callables for any other
+    backend; each receives the executing `AgentThread`, and an optional `preparing=` callable runs
+    before each recall to set up state based on the agent's instructions (it must be idempotent
+    per agent and thread).
+
+Two behavioral notes:
+
+- Context is remembered only when a turn completes and its output stream is fully consumed;
+    turns abandoned mid-stream or failing with an error are not persisted.
+- Memory steps (`prepare_step`, `recall_step`, `remember_step`) require an `AgentThread` bound in
+    the current context and raise `AgentException` otherwise. Agents bind one stamped with their
+    own URI automatically; when running memory steps outside an agent, enter a scope with an
+    `AgentThread` instance first.
 
 ## 3. Preserve Thread And Metadata
 
@@ -232,7 +251,7 @@ agent: Agent = Agent(
 
 async with ctx.scope(
     "agents.context",
-    AgentThread.of(identifier=uuid4(), meta={"source": "outer"}),
+    AgentThread.of(identifier=uuid4(), agent_uri="agent://outer", meta={"source": "outer"}),
 ):
     stream: AsyncIterable[MultimodalContentPart | ProcessingEvent] = agent.call(
         input="hello",
@@ -394,6 +413,9 @@ The public agents API exported from `draive` includes:
 - `AgentExecuting`
 - `AgentIdentity`
 - `AgentMemory`
+- `AgentMemoryPreparing`
+- `AgentMemoryRecalling`
+- `AgentMemoryRemembering`
 - `AgentMessage`
 - `AgentThread`
 - `ProcessingEvent`

docs/guides/Postgres.md

@@ -4,8 +4,8 @@ Draive ships with Postgres-backed implementations for common persistence interfa
 relational storage into your workflows without writing adapters. All helpers live in
 `draive.postgres` and reuse the shared `haiway.postgres.Postgres` connection states.
 
-Current adapters include `PostgresConfigurationRepository`, `PostgresTemplatesRepository`, and
-`PostgresVectorIndex`.
+Current adapters include `PostgresConfigurationRepository`, `PostgresTemplatesRepository`,
+`PostgresVectorIndex`, and `PostgresAgentMemory`.
 
 ## Bootstrapping the Postgres context
 
@@ -211,6 +211,49 @@ payload JSON. Requirements are translated to SQL expressions (for example,
 `AttributeRequirement.equal` becomes `payload #>> '{text}' = $2`). Unsupported operators raise
 `NotImplementedError`, ensuring the query surface remains explicit.
 
+## AgentMemory implementation
+
+`PostgresAgentMemory` persists agent conversation context across turns, keyed by the executing
+agent URI and the conversation thread identifier - both carried on the `AgentThread` passed to
+each memory operation. See the [Agents](./Agents.md) guide for how `AgentMemory` participates in
+agent execution.
+
+Run the schema migration once with an acquired connection bound in context:
+
+```python
+from draive import ctx
+from draive.postgres import Postgres, PostgresAgentMemory
+
+async with ctx.scope("migration"):
+    async with Postgres.acquire_connection() as connection:
+        with ctx.updating(connection):
+            await PostgresAgentMemory.migrate()
+```
+
+Then create a memory instance and pass it to an agent:
+
+```python
+from draive import Agent
+from draive.postgres import PostgresAgentMemory
+
+assistant = Agent.generative(
+    "assistant",
+    instructions="You are a concise support assistant.",
+    memory=PostgresAgentMemory.instance(),
+)
+```
+
+Storage semantics:
+
+- Context is stored as immutable snapshots: every remember inserts the full context as a new row
+    and recall reads back the latest one whole, so steps may compact, summarize, or replace the
+    context freely between turns.
+- Persistence is lock-free and write-only: previous snapshots are never modified or deleted and
+    remain available for tracking and verification. Snapshot history grows without bound over the
+    lifetime of a thread.
+- A single memory instance can serve multiple agents; recalled and remembered context is isolated
+    per agent URI and thread, both taken from the executing `AgentThread` at call time.
+
 ## Putting it together
 
 Combine these adapters with higher-level Draive components to centralise operational data in

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "uv_build"
 [project]
 name = "draive"
 description = "Framework designed to simplify and accelerate the development of LLM-based applications."
-version = "0.113.1"
+version = "0.114.0"
 readme = "README.md"
 maintainers = [
   { name = "Kacper Kaliński", email = "kacper.kalinski@miquido.com" },

src/draive/agents/__init__.py

@@ -5,6 +5,9 @@
     AgentException,
     AgentExecuting,
     AgentIdentity,
+    AgentMemoryPreparing,
+    AgentMemoryRecalling,
+    AgentMemoryRemembering,
     AgentMessage,
     AgentThread,
     AgentUnavailable,
@@ -16,6 +19,9 @@
     "AgentExecuting",
     "AgentIdentity",
     "AgentMemory",
+    "AgentMemoryPreparing",
+    "AgentMemoryRecalling",
+    "AgentMemoryRemembering",
     "AgentMessage",
     "AgentThread",
     "AgentUnavailable",