Stateful inference for Pydantic AI. Add snapshot-based state management to any existing Pydantic AI app without changing a line of application code.
Every LLM conversation today is stateless. Your application sends the entire message history on every turn. The model re-reads everything from scratch. You pay for all those tokens again, and latency grows with every message.
On a 50-turn conversation, a stateless system processes approximately 25,000 tokens total. Most of that is redundant. Standard alternatives require either major application rewrites or inference server fork maintenance.
The Engram Pydantic AI Provider wraps Engram's snapshot-based persistent state system as a Pydantic AI provider. The provider automatically detects redundant context in your messages, restores a saved model state snapshot (approximately 2ms), and only processes the new tokens.
This works as a drop-in replacement. Stateless Pydantic AI apps get stateful benefits with zero code changes. On a 50-turn conversation, prefix detection reduces token processing by 93.8%, bringing the total from 25,000 tokens down to 1,550 tokens.
For applications that want explicit control, the provider also exposes all 5 Engram snapshot endpoints (save, restore, list, info, delete) as first-class methods.
Token savings: 93.8% reduction on repeated conversation prefixes (tested on 50-turn interaction: 25,000 tokens reduced to 1,550).
Restore latency: 2ms average for model state restore using Engram's snapshot persistence.
Test coverage: 39 passing tests covering provider initialization, state tracking, prefix matching, and error handling.
Validated models: IBM Granite, NVIDIA Nemotron, Alibaba Qwen3, Codestral, and hybrid Mamba2 architectures.
| Model | Architecture | Status |
|---|---|---|
| IBM Granite 4.0-H-tiny (4B) | Mamba2 + Attention hybrid | PASS |
| IBM Granite 4.0-H-small (8B) | Mamba2 + Attention hybrid | PASS |
| NVIDIA Nemotron-Cascade-2-30B | MoE Mamba2 hybrid | PASS |
| Alibaba Qwen3-Next-80B | Mamba2 + Attn + MoE | PASS |
| Codestral 7B | Pure Mamba2 | PASS |
| NVIDIA Nemotron-3-Super-120B FP8 | LatentMoE Mamba2 hybrid | BLOCKED (SM89+) |
pip install engram-pydantic-providerOr from GitHub:
pip install git+https://github.com/Clarit-AI/Engram-Pydantic-Provider.gitWith tokenizer support for accurate token estimates:
pip install engram-pydantic-provider[transformers]# Install Engram
pip install -e "python/" # from the Engram repo
# Start the server with snapshot persistence
python -m sglang.launch_server \
--model-path ibm-granite/granite-4.0-h-tiny \
--enable-snapshot-persistence \
--snapshot-dir ./snapshots \
--mamba-scheduler-strategy no_buffer \
--disable-radix-cache \
--port 30000from pydantic_ai import Agent
from engram_pydantic_provider import EngramModel
# Drop-in replacement - zero code changes for stateless apps
agent = Agent(EngramModel("granite-4.0-h-tiny"))
result = await agent.run("Hello!")| Variable | Default | Description |
|---|---|---|
ENGRAM_BASE_URL |
http://localhost:30000 |
Engram server URL |
ENGRAM_API_KEY |
None |
API key for authenticated servers |
ENGRAM_ADMIN_API_KEY |
None |
API key for snapshot routes |
ENGRAM_AUTO_SAVE |
true |
Save snapshot after every turn |
ENGRAM_STATEFUL_MODE |
auto |
Default operating mode: auto, stateless, or explicit |
ENGRAM_TOKENIZER_PATH |
None |
Local tokenizer path for air-gapped deployments |
ENGRAM_SNAPSHOT_TIMEOUT |
30.0 |
HTTP timeout in seconds for snapshot operations |
from pydantic_ai import Agent
from engram_pydantic_provider import EngramModel, EngramConfig
config = EngramConfig(
base_url="https://your-engram-instance.com",
api_key="your-api-key",
stateful_mode="auto", # or "stateless" or "explicit"
)
agent = Agent(EngramModel("granite-4.0-h-tiny", config=config))The provider automatically:
- Generates a pseudo-ID from the first 2 messages (or uses your
conversation_id) - On subsequent calls, detects if the incoming messages are a prefix match with known state
- If matched, restores the snapshot (2ms) and sends only new messages for generation
- Saves a new snapshot after the response
This is the zero-config mode. Stateless applications get stateful benefits automatically.
result = await agent.run("Continue the conversation")
# Provider auto-detects prefix, restores snapshot, processes only new tokensStandard OpenAI-compatible pass-through. No snapshots, no state tracking. Useful for testing or bypassing the stateful system.
result = await agent.run(
"Stateless request",
model_settings={"engram_mode": "stateless"}
)You manage snapshots directly via conversation_id and restore targets. The provider does not auto-detect prefixes. It only restores when you specify it.
result = await agent.run(
"Explicit session",
model_settings={
"engram_mode": "explicit",
"engram_conversation_id": "my-session",
"engram_restore_turn": 5,
}
)┌─────────────────────────────────────────────────────┐
│ Your Application │
│ (Pydantic AI Agent) │
└─────────────────────┬───────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────┐
│ EngramModel │
│ (extends OpenAIChatModel) │
│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │
│ │ Tracker │ │ ContextDiff │ │ SnapshotCli │ │
│ │ (conv IDs) │ │ (prefix) │ │ (5 methods) │ │
│ └─────────────┘ └─────────────┘ └─────────────┘ │
│ ┌─────────────┐ ┌─────────────┐ │
│ │ Tokenizer │ │ EngramModel │ │
│ │ (encoding) │ │ (override) │ │
│ └─────────────┘ └─────────────┘ │
└─────────────────────┬───────────────────────────────┘
│
┌───────────┴───────────┐
▼ ▼
┌─────────────────┐ ┌─────────────────┐
│restore_snapshot │ │ /v1/chat/ │
│ (2ms) │ │ completions │
└────────┬────────┘ └────────┬────────┘
│ │
└───────────┬───────────┘
▼
┌─────────────────┐
│ save_snapshot │
│ (auto or manual)│
└─────────────────┘
│
▼
┌─────────────────────────────────────────────────────┐
│ Engram Server │
│ (SGLang + Mamba + Snapshot Persistence) │
└─────────────────────────────────────────────────────┘
| Component | File | Purpose |
|---|---|---|
EngramModel |
model.py |
Main model class, extends OpenAIChatModel with snapshot lifecycle |
EngramProvider |
provider.py |
Provider for Engram (stateful SGLang) server |
ConversationTracker |
tracker.py |
Thread-safe conversation state tracking with LRU eviction |
ContextDiffer |
differ.py |
Prefix matching algorithm (strips 1 to 5 messages from tail) |
SnapshotClient |
snapshot.py |
Async HTTP client for all 5 Engram snapshot endpoints |
TokenizerWrapper |
tokenizer.py |
Lazy-loaded tokenizer for continuation encoding |
A typical multi-turn chat application sends the full message history on every turn:
Turn 1: [system, user1] 50 tokens
Turn 2: [system, user1, asst1, user2] 120 tokens
Turn 3: [system, user1, asst1, user2, asst2, user3] 210 tokens
...
Turn 20: [system, user1, ..., user20] 5,000 tokens
Total tokens across 20 turns: approximately 25,000 tokens. Most of that is redundant re-processing.
The provider detects that each new request extends a known conversation:
Turn 1: Process [system, user1] (50 tokens), save snapshot
Turn 2: Restore snapshot (2ms), process [asst1, user2] (70 tokens), save
Turn 3: Restore snapshot (2ms), process [asst2, user3] (90 tokens), save
...
Turn 20: Restore snapshot (2ms), process [asst19, user20] (250 tokens)
Total tokens processed: approximately 1,550 tokens, a 93.8% reduction.
- Hash the stored message history for each known conversation
- On each incoming request, strip 1 to 5 messages from the tail
- Hash each stripped prefix and compare against known hashes
- If a match is found, restore that snapshot and process only the delta messages
- If no match, run a standard full-prefill completion
When no conversation_id is provided, the provider generates a deterministic pseudo-ID from the first 2 messages (typically system prompt and first user message). This scopes prefix matching to prevent cross-conversation contamination.
The provider handles streaming with restore-before and save-after semantics:
- Restore synchronously before the stream starts (2ms)
- Proxy SSE chunks through to your application
- Save asynchronously after the stream completes (fire-and-forget)
If your application cancels the stream mid-generation, the save is skipped.
For applications that need direct snapshot control:
from engram_pydantic_provider import SnapshotClient, EngramConfig
config = EngramConfig(base_url="http://localhost:30000")
client = SnapshotClient(base_url=config.base_url, auth_key=config.snapshot_auth_key)
# Save
save_result = await client.save_snapshot(
rid="request-id-from-response",
conversation_id="session-abc",
turn_number=5,
)
# List all snapshots for a conversation
snapshots = await client.list_snapshots("session-abc")
for s in snapshots:
print(f"Turn {s.turn_number}: {s.snapshot_id}")
# Get info on a specific snapshot
info = await client.get_snapshot_info("session-abc", turn_number=3)
# Restore
restore_result = await client.restore_snapshot(
conversation_id="session-abc",
turn_number=3,
create_new_request=True,
max_new_tokens=256,
)
# Delete
await client.delete_snapshot("session-abc", turn_number=1)| Failure | Behavior |
|---|---|
| Restore HTTP error (4xx/5xx) | Log warning, fall back to full prefill |
| Snapshot not found (404) | Same as above; fallback to full prefill |
| Save failure after completion | Log warning silently, do not surface to caller |
| Tokenizer load failure | Skip Path B optimization, fall back to restore-only |
| Prefix match failure | Fall back to standard generation |
The ConversationTracker is process-local. In multi-worker deployments, each worker maintains independent state. Cross-worker requests will silently fall back to full prefill.
Recommendations for distributed deployments:
- Always provide an explicit
conversation_idviamodel_settings - Use
stateful_mode: "explicit"for predictable behavior
# Install test dependencies
pip install pytest pytest-asyncio
# Run all tests (no GPU required)
python -m pytest tests/ -vAll 39 tests run without GPU.
Tests cover:
- Provider initialization and configuration
- Thread-safe conversation state tracking and LRU eviction
- Prefix matching algorithm and delta computation
- SnapshotClient CRUD operations and error handling
- Model profile presets (minimal, standard, full)
- Mode resolution (auto, stateless, explicit)
Engram Pydantic AI Provider is part of the Clarit.AI open-source ecosystem. Engram Pydantic AI Provider focuses on drop-in stateful inference for existing Pydantic AI applications, while Engram focuses on foundational Mamba serving with snapshot persistence.
Contributions are welcome. The provider follows the same structure as Pydantic AI upstream projects.
Areas for contribution:
- Additional model architecture validation
- Distributed state tracking backends (Redis, etc.)
- Performance benchmarking and optimization
Start by opening an issue to discuss your idea, then submit a pull request with tests for your changes.
MIT, same as Pydantic AI.
Built on Engram by Clarit.AI, which extends SGLang with persistent Mamba state management. Designed for the Pydantic AI ecosystem by Pydantic.
The Mamba architecture was developed by Albert Gu and Tri Dao. Mamba2 was developed by Tri Dao and Albert Gu at Carnegie Mellon University.
- Engram: https://github.com/Clarit-AI/Engram
- Pydantic AI: https://github.com/pydantic/pydantic-ai
- Engram API Guide: https://github.com/Clarit-AI/Engram/blob/main/docs/stateful_mamba/api_guide.md
- Clarit.AI: https://clarit.ai