langgraph-dynamodb-store 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- langgraph_dynamodb_store-0.1.0/.github/workflows/publish.yml +45 -0
- langgraph_dynamodb_store-0.1.0/.github/workflows/test.yml +27 -0
- langgraph_dynamodb_store-0.1.0/.gitignore +10 -0
- langgraph_dynamodb_store-0.1.0/LICENSE +21 -0
- langgraph_dynamodb_store-0.1.0/MEMORY_LAYER.md +178 -0
- langgraph_dynamodb_store-0.1.0/PKG-INFO +204 -0
- langgraph_dynamodb_store-0.1.0/README.md +173 -0
- langgraph_dynamodb_store-0.1.0/pyproject.toml +56 -0
- langgraph_dynamodb_store-0.1.0/src/memory_layer/__init__.py +4 -0
- langgraph_dynamodb_store-0.1.0/src/memory_layer/retrieval.py +35 -0
- langgraph_dynamodb_store-0.1.0/src/memory_layer/store.py +159 -0
- langgraph_dynamodb_store-0.1.0/src/memory_layer/types.py +15 -0
- langgraph_dynamodb_store-0.1.0/src/memory_layer/writer.py +75 -0
- langgraph_dynamodb_store-0.1.0/tests/conftest.py +41 -0
- langgraph_dynamodb_store-0.1.0/tests/test_store.py +128 -0
- langgraph_dynamodb_store-0.1.0/tests/test_writer.py +86 -0
|
@@ -0,0 +1,45 @@
|
|
|
1
|
+
name: Publish to PyPI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
release:
|
|
5
|
+
types: [published]
|
|
6
|
+
|
|
7
|
+
jobs:
|
|
8
|
+
build:
|
|
9
|
+
runs-on: ubuntu-latest
|
|
10
|
+
steps:
|
|
11
|
+
- uses: actions/checkout@v4
|
|
12
|
+
|
|
13
|
+
- name: Set up Python
|
|
14
|
+
uses: actions/setup-python@v5
|
|
15
|
+
with:
|
|
16
|
+
python-version: "3.12"
|
|
17
|
+
|
|
18
|
+
- name: Build sdist and wheel
|
|
19
|
+
run: |
|
|
20
|
+
pip install build
|
|
21
|
+
python -m build
|
|
22
|
+
|
|
23
|
+
- name: Upload build artifacts
|
|
24
|
+
uses: actions/upload-artifact@v4
|
|
25
|
+
with:
|
|
26
|
+
name: dist
|
|
27
|
+
path: dist/
|
|
28
|
+
|
|
29
|
+
publish:
|
|
30
|
+
needs: build
|
|
31
|
+
runs-on: ubuntu-latest
|
|
32
|
+
# Trusted Publishing (OIDC) — no PyPI API token/secret needed. Requires a
|
|
33
|
+
# pending publisher registered at https://pypi.org/manage/account/publishing/
|
|
34
|
+
# pointing at this repo + this exact workflow filename before the first release.
|
|
35
|
+
permissions:
|
|
36
|
+
id-token: write
|
|
37
|
+
steps:
|
|
38
|
+
- name: Download build artifacts
|
|
39
|
+
uses: actions/download-artifact@v4
|
|
40
|
+
with:
|
|
41
|
+
name: dist
|
|
42
|
+
path: dist/
|
|
43
|
+
|
|
44
|
+
- name: Publish to PyPI
|
|
45
|
+
uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
name: Tests
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches: [main]
|
|
6
|
+
pull_request:
|
|
7
|
+
branches: [main]
|
|
8
|
+
|
|
9
|
+
jobs:
|
|
10
|
+
test:
|
|
11
|
+
runs-on: ubuntu-latest
|
|
12
|
+
strategy:
|
|
13
|
+
matrix:
|
|
14
|
+
python-version: ["3.11", "3.12"]
|
|
15
|
+
steps:
|
|
16
|
+
- uses: actions/checkout@v4
|
|
17
|
+
|
|
18
|
+
- name: Set up Python ${{ matrix.python-version }}
|
|
19
|
+
uses: actions/setup-python@v5
|
|
20
|
+
with:
|
|
21
|
+
python-version: ${{ matrix.python-version }}
|
|
22
|
+
|
|
23
|
+
- name: Install package with dev dependencies
|
|
24
|
+
run: pip install -e ".[dev]"
|
|
25
|
+
|
|
26
|
+
- name: Run tests
|
|
27
|
+
run: pytest -v
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Mauricio Neira
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,178 @@
|
|
|
1
|
+
# Memory Layer — Design Document
|
|
2
|
+
|
|
3
|
+
## Visión
|
|
4
|
+
|
|
5
|
+
Biblioteca Python que provee **memoria cross-thread** para agentes LangGraph. Implementa el protocolo `BaseStore` de LangGraph para que cualquier grafo pueda leer y escribir memorias de usuario sin acoplar lógica de storage al agente.
|
|
6
|
+
|
|
7
|
+
No es una feature de Atlas — es una dependencia que Atlas (y cualquier otro agente LangGraph) puede importar.
|
|
8
|
+
|
|
9
|
+
---
|
|
10
|
+
|
|
11
|
+
## Problema
|
|
12
|
+
|
|
13
|
+
LangGraph `checkpointer` persiste el estado **dentro** de un thread (conversación). Cuando el usuario abre una nueva sesión, el grafo empieza de cero — no sabe que el usuario prefiere análisis por canal, que trabaja con la cuenta de Nike, o que la última vez quedó una pregunta sin resolver.
|
|
14
|
+
|
|
15
|
+
El `store` de LangGraph resuelve esto: memoria **cross-thread**, compartida entre todas las sesiones de un mismo usuario.
|
|
16
|
+
|
|
17
|
+
---
|
|
18
|
+
|
|
19
|
+
## Tipos de Memoria
|
|
20
|
+
|
|
21
|
+
| Tipo | Descripción | Ejemplo |
|
|
22
|
+
|---|---|---|
|
|
23
|
+
| `semantic` | Preferencias y hechos estables del usuario | "Prefiere respuestas en español", "trabaja con cuenta Nike" |
|
|
24
|
+
| `episodic` | Eventos pasados relevantes | "El 2025-05-10 analizamos fatiga de creativos Q4 de Nike" |
|
|
25
|
+
| `procedural` | Patrones de trabajo recurrentes | "Siempre empieza con breakdown por canal, luego creativos" |
|
|
26
|
+
|
|
27
|
+
---
|
|
28
|
+
|
|
29
|
+
## Scopes
|
|
30
|
+
|
|
31
|
+
```
|
|
32
|
+
user:{cognito_sub} # preferencias personales — privado al usuario
|
|
33
|
+
instance:{instance_id} # contexto compartido del tenant (todos los usuarios lo ven)
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
El scope se mapea directamente al `namespace` de LangGraph Store: `("user", user_id)`.
|
|
37
|
+
|
|
38
|
+
---
|
|
39
|
+
|
|
40
|
+
## Integración con LangGraph
|
|
41
|
+
|
|
42
|
+
LangGraph compila el grafo con `store` al lado del `checkpointer`:
|
|
43
|
+
|
|
44
|
+
```python
|
|
45
|
+
from memory_layer import DynamoDBStore
|
|
46
|
+
|
|
47
|
+
store = DynamoDBStore(table_name="memory-layer-dev-memories")
|
|
48
|
+
graph = builder.compile(checkpointer=checkpointer, store=store)
|
|
49
|
+
```
|
|
50
|
+
|
|
51
|
+
Dentro de un nodo, el store se accede vía `RunnableConfig`:
|
|
52
|
+
|
|
53
|
+
```python
|
|
54
|
+
from langgraph.store.base import BaseStore
|
|
55
|
+
|
|
56
|
+
def supervisor_node(state: AtlasState, config: RunnableConfig, store: BaseStore) -> dict:
|
|
57
|
+
user_id = state["context"].get("user_id", "")
|
|
58
|
+
memories = store.search(("user", user_id), query=state["messages"][-1].content, limit=5)
|
|
59
|
+
# inject relevant memories into system prompt
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
Al final de la conversación, un nodo `memory_writer` opcional extrae y persiste hechos nuevos:
|
|
63
|
+
|
|
64
|
+
```python
|
|
65
|
+
def memory_writer_node(state: AtlasState, store: BaseStore) -> dict:
|
|
66
|
+
# run LLM to extract key facts from the conversation
|
|
67
|
+
# store.put(("user", user_id), memory_id, {"content": ..., "type": "semantic"})
|
|
68
|
+
return {}
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
---
|
|
72
|
+
|
|
73
|
+
## Estrategias de Retrieval
|
|
74
|
+
|
|
75
|
+
### 1. Simple (MVP)
|
|
76
|
+
Fetch N memorias más recientes del usuario, inyectar todas en el system prompt.
|
|
77
|
+
- Bueno para: pocos usuarios, pocas memorias por usuario
|
|
78
|
+
- Límite: el context window crece con las memorias
|
|
79
|
+
|
|
80
|
+
### 2. Semántica (Phase 2)
|
|
81
|
+
Embed la query del usuario + las memorias. Cosine similarity para seleccionar las más relevantes.
|
|
82
|
+
- Requiere: embeddings (OpenAI `text-embedding-3-small` o similar)
|
|
83
|
+
- Storage: vector en DynamoDB (como string serializado) o migrar a pgvector/Pinecone
|
|
84
|
+
|
|
85
|
+
### 3. Por tipo (Phase 2)
|
|
86
|
+
Filtrar por `type` antes de rankear — siempre incluir `semantic` (preferencias), limitar `episodic`.
|
|
87
|
+
|
|
88
|
+
**Arrancar con Simple. Migrar a Semántica cuando el número de memorias por usuario supere ~20.**
|
|
89
|
+
|
|
90
|
+
---
|
|
91
|
+
|
|
92
|
+
## Estrategias de Escritura
|
|
93
|
+
|
|
94
|
+
### Automática (recomendada)
|
|
95
|
+
Nodo `memory_writer` al final del grafo. LLM extrae hechos nuevos de la conversación y los persiste. Se ejecuta solo si hubo interacción con un agente specialist (no para saludos/clarificaciones).
|
|
96
|
+
|
|
97
|
+
### Explícita
|
|
98
|
+
El agente escribe memorias cuando el usuario dice "recordá que..." o cuando detecta un hecho relevante durante la conversación. Requiere una tool `save_memory` disponible para los agentes.
|
|
99
|
+
|
|
100
|
+
### Ambas
|
|
101
|
+
La combinación más robusta. Explícita para capturar intenciones directas del usuario; automática como red de seguridad.
|
|
102
|
+
|
|
103
|
+
---
|
|
104
|
+
|
|
105
|
+
## Data Model — DynamoDB
|
|
106
|
+
|
|
107
|
+
**Tabla**: `memory-layer-{stage}-memories`
|
|
108
|
+
|
|
109
|
+
| Campo | Tipo | Descripción |
|
|
110
|
+
|---|---|---|
|
|
111
|
+
| `owner_id` | PK (String) | Scope + id: `"user:abc"`, `"instance:bunker"` |
|
|
112
|
+
| `memory_id` | SK (String) | UUID |
|
|
113
|
+
| `type` | String | `"semantic"` \| `"episodic"` \| `"procedural"` |
|
|
114
|
+
| `content` | String | Texto libre de la memoria |
|
|
115
|
+
| `created_at` | String | ISO 8601 UTC |
|
|
116
|
+
| `expires_at` | Number | TTL Unix timestamp (opcional — memorias semánticas no expiran) |
|
|
117
|
+
| `embedding` | String | JSON serializado del vector (Phase 2) |
|
|
118
|
+
| `session_id` | String | Thread que originó esta memoria (trazabilidad) |
|
|
119
|
+
|
|
120
|
+
**GSI**: `owner_id-created_at-index` para fetch por recencia.
|
|
121
|
+
|
|
122
|
+
---
|
|
123
|
+
|
|
124
|
+
## Arquitectura del Proyecto
|
|
125
|
+
|
|
126
|
+
```
|
|
127
|
+
memory-layer/
|
|
128
|
+
├── src/
|
|
129
|
+
│ └── memory_layer/
|
|
130
|
+
│ ├── __init__.py
|
|
131
|
+
│ ├── store.py # DynamoDBStore — implementa LangGraph BaseStore
|
|
132
|
+
│ ├── types.py # MemoryRecord, MemoryType, MemoryScope TypedDicts
|
|
133
|
+
│ ├── retrieval.py # Estrategias: SimpleRetrieval, SemanticRetrieval
|
|
134
|
+
│ └── writer.py # MemoryWriter — extracción LLM de hechos
|
|
135
|
+
├── tests/
|
|
136
|
+
│ ├── conftest.py
|
|
137
|
+
│ ├── test_store.py
|
|
138
|
+
│ └── test_writer.py
|
|
139
|
+
├── pyproject.toml
|
|
140
|
+
└── MEMORY_LAYER.md
|
|
141
|
+
```
|
|
142
|
+
|
|
143
|
+
---
|
|
144
|
+
|
|
145
|
+
## Plan de Implementación
|
|
146
|
+
|
|
147
|
+
### Phase 1 — Store + Simple Retrieval
|
|
148
|
+
- [ ] `DynamoDBStore` implementando `langgraph.store.base.BaseStore`
|
|
149
|
+
- [ ] `SimpleRetrieval` — fetch N más recientes, sin embeddings
|
|
150
|
+
- [ ] `MemoryRecord` TypedDict + validación
|
|
151
|
+
- [ ] Tests con DynamoDB Local
|
|
152
|
+
|
|
153
|
+
### Phase 2 — Semantic Retrieval
|
|
154
|
+
- [ ] Embeddings via `langchain-openai` (`text-embedding-3-small`)
|
|
155
|
+
- [ ] `SemanticRetrieval` — cosine similarity sobre embeddings en DynamoDB
|
|
156
|
+
- [ ] Migración de registros existentes (backfill embeddings)
|
|
157
|
+
|
|
158
|
+
### Phase 3 — Memory Writer
|
|
159
|
+
- [ ] `MemoryWriter` — nodo LangGraph que extrae hechos al final de cada conversación
|
|
160
|
+
- [ ] Deduplicación — no escribir memorias redundantes (LLM judge o embedding similarity)
|
|
161
|
+
- [ ] Expiración diferenciada por tipo — `episodic` TTL 90 días, `semantic` sin TTL
|
|
162
|
+
|
|
163
|
+
### Integración en Atlas
|
|
164
|
+
- [ ] Importar `memory_layer` como dependencia en `requirements.txt`
|
|
165
|
+
- [ ] Pasar `DynamoDBStore` al compilar el grafo en `core/graph.py`
|
|
166
|
+
- [ ] Nodo `memory_writer` opcional al final del grafo (skill `user_memory`)
|
|
167
|
+
- [ ] Inyección de memorias en `supervisor_node` vía `store.search()`
|
|
168
|
+
|
|
169
|
+
---
|
|
170
|
+
|
|
171
|
+
## Decisiones Pendientes
|
|
172
|
+
|
|
173
|
+
| Decisión | Opciones | Estado |
|
|
174
|
+
|---|---|---|
|
|
175
|
+
| ¿Embeddings propios o pgvector? | DynamoDB + cosine en Python vs pgvector en RDS | Pendiente — depende del volumen de usuarios |
|
|
176
|
+
| ¿Memory writer siempre activo o como skill? | Siempre vs `required_skill="user_memory"` | Pendiente |
|
|
177
|
+
| ¿TTL para memorias semánticas? | Sin TTL vs 1 año | Pendiente |
|
|
178
|
+
| ¿Límite de memorias por usuario? | Sin límite vs max 100 + archivado | Pendiente |
|
|
@@ -0,0 +1,204 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: langgraph-dynamodb-store
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: A DynamoDB-backed BaseStore for LangGraph — cross-thread memory for agents, no Postgres/pgvector required
|
|
5
|
+
Project-URL: Homepage, https://github.com/mauriciosneira/memory-layer
|
|
6
|
+
Project-URL: Issues, https://github.com/mauriciosneira/memory-layer/issues
|
|
7
|
+
Author: Mauricio Neira
|
|
8
|
+
License-Expression: MIT
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Keywords: agents,dynamodb,langchain,langgraph,llm,memory
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
17
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
18
|
+
Requires-Python: >=3.11
|
|
19
|
+
Requires-Dist: boto3>=1.35
|
|
20
|
+
Requires-Dist: langchain-core>=0.3
|
|
21
|
+
Requires-Dist: langgraph>=0.2
|
|
22
|
+
Requires-Dist: pydantic>=2.0
|
|
23
|
+
Provides-Extra: dev
|
|
24
|
+
Requires-Dist: moto[dynamodb]>=5.0; extra == 'dev'
|
|
25
|
+
Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
|
|
26
|
+
Requires-Dist: pytest>=8.0; extra == 'dev'
|
|
27
|
+
Provides-Extra: semantic
|
|
28
|
+
Requires-Dist: langchain-openai>=0.2; extra == 'semantic'
|
|
29
|
+
Requires-Dist: numpy>=1.26; extra == 'semantic'
|
|
30
|
+
Description-Content-Type: text/markdown
|
|
31
|
+
|
|
32
|
+
# memory-layer
|
|
33
|
+
|
|
34
|
+
[](https://github.com/mauriciosneira/memory-layer/actions/workflows/test.yml)
|
|
35
|
+
[](LICENSE)
|
|
36
|
+
|
|
37
|
+
**A DynamoDB-backed `BaseStore` for LangGraph — cross-thread memory for your agents without standing up Postgres/pgvector.**
|
|
38
|
+
|
|
39
|
+
LangGraph's checkpointer persists state *inside* a thread — when a user opens a new conversation, the graph starts from zero. The `store` protocol is LangGraph's answer to that: memory that survives *across* threads, shared by every session for the same user (or the same tenant). LangGraph ships an official store for Postgres. If your stack is already DynamoDB — which a lot of serverless/Fargate deployments are — there wasn't an official option. `memory-layer` is that option.
|
|
40
|
+
|
|
41
|
+
```python
|
|
42
|
+
from memory_layer import DynamoDBStore
|
|
43
|
+
|
|
44
|
+
store = DynamoDBStore(table_name="my-app-memories")
|
|
45
|
+
graph = builder.compile(checkpointer=checkpointer, store=store)
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
That's the whole integration. No new infra beyond one DynamoDB table you probably already know how to provision.
|
|
49
|
+
|
|
50
|
+
---
|
|
51
|
+
|
|
52
|
+
## Why this exists
|
|
53
|
+
|
|
54
|
+
- **You're already on DynamoDB.** Adding Postgres + pgvector just for agent memory is a real infra cost — a new engine, a new backup story, a new thing to monitor — for a feature that, for most products, doesn't need vector search on day one.
|
|
55
|
+
- **LangGraph's `store` protocol is a clean seam.** It's designed so storage is swappable — your agent code shouldn't care whether memories live in Postgres, Dynamo, or Redis. This fills the Dynamo gap in that seam.
|
|
56
|
+
- **Memory doesn't have to mean embeddings.** Most products get real value from "the last N things we know about this user," fetched by recency — no vector index required. `memory-layer` starts there, and gives you a clean place to add semantic ranking later if you actually need it.
|
|
57
|
+
|
|
58
|
+
## Features
|
|
59
|
+
|
|
60
|
+
- **`DynamoDBStore`** — a complete `langgraph.store.base.BaseStore` implementation: `get`/`put`/`search` (and their async counterparts), real pagination via DynamoDB's `LastEvaluatedKey` (not a "hope the first page has enough" heuristic), per-item filtering, and per-type TTL (e.g. auto-expire `episodic` memories after 90 days while `semantic` ones never expire).
|
|
61
|
+
- **`SimpleRetrieval`** — fetch a user's N most recent memories and turn them into a ready-to-inject prompt block. No embeddings, no extra dependencies.
|
|
62
|
+
- **`MemoryWriter`** — an LLM-driven extraction step: hand it a conversation, it classifies and persists the facts worth remembering, via `with_structured_output` (a real schema-validated response, not a hand-rolled JSON parser hoping the model didn't wrap the array in a sentence).
|
|
63
|
+
- **Scopes, not just users.** Namespaces are plain tuples (`("user", user_id)`, `("instance", tenant_id)`) — model per-user memory, per-tenant shared context, or your own scope, however your product actually shapes ownership.
|
|
64
|
+
|
|
65
|
+
## Install
|
|
66
|
+
|
|
67
|
+
```bash
|
|
68
|
+
pip install langgraph-dynamodb-store
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
Want LLM-driven extraction? `MemoryWriter` takes any LangChain `BaseChatModel` — bring the one you already use, no extra install needed. `numpy`/`langchain-openai` are only required for the semantic-retrieval extra (see [Roadmap](#roadmap)):
|
|
72
|
+
|
|
73
|
+
```bash
|
|
74
|
+
pip install "langgraph-dynamodb-store[semantic]"
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
## DynamoDB table
|
|
78
|
+
|
|
79
|
+
One table, one GSI. Create it however you provision infra (CDK/Terraform/console) — here's the raw shape via the AWS CLI, if you just want to try it out:
|
|
80
|
+
|
|
81
|
+
```bash
|
|
82
|
+
aws dynamodb create-table \
|
|
83
|
+
--table-name my-app-memories \
|
|
84
|
+
--attribute-definitions \
|
|
85
|
+
AttributeName=owner_id,AttributeType=S \
|
|
86
|
+
AttributeName=memory_id,AttributeType=S \
|
|
87
|
+
AttributeName=created_at,AttributeType=S \
|
|
88
|
+
--key-schema \
|
|
89
|
+
AttributeName=owner_id,KeyType=HASH \
|
|
90
|
+
AttributeName=memory_id,KeyType=RANGE \
|
|
91
|
+
--global-secondary-indexes '[{
|
|
92
|
+
"IndexName": "owner_id-created_at-index",
|
|
93
|
+
"KeySchema": [
|
|
94
|
+
{"AttributeName": "owner_id", "KeyType": "HASH"},
|
|
95
|
+
{"AttributeName": "created_at", "KeyType": "RANGE"}
|
|
96
|
+
],
|
|
97
|
+
"Projection": {"ProjectionType": "ALL"}
|
|
98
|
+
}]' \
|
|
99
|
+
--billing-mode PAY_PER_REQUEST
|
|
100
|
+
|
|
101
|
+
# Optional but recommended — lets episodic memories actually expire instead of
|
|
102
|
+
# accumulating forever. memory-layer sets the `ttl` attribute; DynamoDB does the rest.
|
|
103
|
+
aws dynamodb update-time-to-live \
|
|
104
|
+
--table-name my-app-memories \
|
|
105
|
+
--time-to-live-specification "Enabled=true, AttributeName=ttl"
|
|
106
|
+
```
|
|
107
|
+
|
|
108
|
+
Set `MEMORY_TABLE=my-app-memories` or pass `table_name` explicitly — `DynamoDBStore(table_name="my-app-memories")`.
|
|
109
|
+
|
|
110
|
+
## Quickstart
|
|
111
|
+
|
|
112
|
+
```python
|
|
113
|
+
from memory_layer import DynamoDBStore
|
|
114
|
+
|
|
115
|
+
store = DynamoDBStore(table_name="my-app-memories")
|
|
116
|
+
|
|
117
|
+
namespace = ("user", "user-123")
|
|
118
|
+
|
|
119
|
+
store.put(namespace, "mem-1", {"content": "Prefers responses in Spanish", "type": "semantic"})
|
|
120
|
+
store.put(namespace, "mem-2", {"content": "Reviewed Q3 numbers on 2026-08-01", "type": "episodic"})
|
|
121
|
+
|
|
122
|
+
memories = store.search(namespace, limit=10)
|
|
123
|
+
for m in memories:
|
|
124
|
+
print(m.value["content"])
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
### Inside a LangGraph node
|
|
128
|
+
|
|
129
|
+
LangGraph injects `store` into any node whose signature asks for it:
|
|
130
|
+
|
|
131
|
+
```python
|
|
132
|
+
from langgraph.store.base import BaseStore
|
|
133
|
+
from langchain_core.runnables import RunnableConfig
|
|
134
|
+
|
|
135
|
+
from memory_layer import DynamoDBStore
|
|
136
|
+
from memory_layer.retrieval import SimpleRetrieval
|
|
137
|
+
|
|
138
|
+
store = DynamoDBStore(table_name="my-app-memories")
|
|
139
|
+
graph = builder.compile(checkpointer=checkpointer, store=store)
|
|
140
|
+
|
|
141
|
+
def supervisor_node(state: AgentState, config: RunnableConfig, store: BaseStore) -> dict:
|
|
142
|
+
user_id = state["context"]["user_id"]
|
|
143
|
+
retrieval = SimpleRetrieval(store, limit=5)
|
|
144
|
+
memories = retrieval.fetch(("user", user_id))
|
|
145
|
+
memory_block = retrieval.to_prompt_block(memories)
|
|
146
|
+
# ...inject memory_block into the system prompt
|
|
147
|
+
return {}
|
|
148
|
+
```
|
|
149
|
+
|
|
150
|
+
### Writing memories back
|
|
151
|
+
|
|
152
|
+
```python
|
|
153
|
+
from memory_layer.writer import MemoryWriter
|
|
154
|
+
|
|
155
|
+
writer = MemoryWriter(llm=your_chat_model, store=store)
|
|
156
|
+
|
|
157
|
+
async def memory_writer_node(state: AgentState) -> dict:
|
|
158
|
+
await writer.extract_and_save(
|
|
159
|
+
namespace=("user", state["context"]["user_id"]),
|
|
160
|
+
messages=state["messages"],
|
|
161
|
+
session_id=state["context"]["session_id"],
|
|
162
|
+
)
|
|
163
|
+
return {}
|
|
164
|
+
```
|
|
165
|
+
|
|
166
|
+
The extraction LLM only needs `with_structured_output` support — every major provider's LangChain integration has it.
|
|
167
|
+
|
|
168
|
+
## Memory types
|
|
169
|
+
|
|
170
|
+
| Type | What it's for | Default TTL |
|
|
171
|
+
|---|---|---|
|
|
172
|
+
| `semantic` | Stable preferences and facts ("prefers Spanish", "works on the Acme account") | none |
|
|
173
|
+
| `episodic` | Specific past events ("reviewed Q3 numbers on 2026-08-01") | 90 days |
|
|
174
|
+
| `procedural` | Recurring work patterns ("always starts with a channel breakdown") | none |
|
|
175
|
+
|
|
176
|
+
TTL policy per type lives in `memory_layer.store._TTL_SECONDS_BY_TYPE` — override it if 90 days isn't the right default for your product.
|
|
177
|
+
|
|
178
|
+
## Scopes
|
|
179
|
+
|
|
180
|
+
A namespace is just a tuple — `memory-layer` doesn't prescribe what it means, but the common shapes are:
|
|
181
|
+
|
|
182
|
+
```python
|
|
183
|
+
("user", cognito_sub) # private to one user
|
|
184
|
+
("instance", tenant_id) # shared across every user of one tenant
|
|
185
|
+
```
|
|
186
|
+
|
|
187
|
+
## Roadmap
|
|
188
|
+
|
|
189
|
+
- **Semantic retrieval** — embed memories + query, rank by cosine similarity, for products that outgrow "most recent N" (roughly ~20+ memories per user is where this starts to matter). Lives behind the `semantic` extra so the core library stays dependency-light.
|
|
190
|
+
- **Deduplication on write** — skip persisting a fact that's a near-duplicate of one already stored.
|
|
191
|
+
- **Bring-your-own embeddings backend** — DynamoDB-native cosine similarity to start; pluggable enough to swap in pgvector/a real vector store later if volume ever justifies it.
|
|
192
|
+
|
|
193
|
+
## Development
|
|
194
|
+
|
|
195
|
+
```bash
|
|
196
|
+
pip install -e ".[dev]"
|
|
197
|
+
pytest
|
|
198
|
+
```
|
|
199
|
+
|
|
200
|
+
Tests run against [`moto`](https://github.com/getmoto/moto) — no real AWS account or network access required.
|
|
201
|
+
|
|
202
|
+
## License
|
|
203
|
+
|
|
204
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,173 @@
|
|
|
1
|
+
# memory-layer
|
|
2
|
+
|
|
3
|
+
[](https://github.com/mauriciosneira/memory-layer/actions/workflows/test.yml)
|
|
4
|
+
[](LICENSE)
|
|
5
|
+
|
|
6
|
+
**A DynamoDB-backed `BaseStore` for LangGraph — cross-thread memory for your agents without standing up Postgres/pgvector.**
|
|
7
|
+
|
|
8
|
+
LangGraph's checkpointer persists state *inside* a thread — when a user opens a new conversation, the graph starts from zero. The `store` protocol is LangGraph's answer to that: memory that survives *across* threads, shared by every session for the same user (or the same tenant). LangGraph ships an official store for Postgres. If your stack is already DynamoDB — which a lot of serverless/Fargate deployments are — there wasn't an official option. `memory-layer` is that option.
|
|
9
|
+
|
|
10
|
+
```python
|
|
11
|
+
from memory_layer import DynamoDBStore
|
|
12
|
+
|
|
13
|
+
store = DynamoDBStore(table_name="my-app-memories")
|
|
14
|
+
graph = builder.compile(checkpointer=checkpointer, store=store)
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
That's the whole integration. No new infra beyond one DynamoDB table you probably already know how to provision.
|
|
18
|
+
|
|
19
|
+
---
|
|
20
|
+
|
|
21
|
+
## Why this exists
|
|
22
|
+
|
|
23
|
+
- **You're already on DynamoDB.** Adding Postgres + pgvector just for agent memory is a real infra cost — a new engine, a new backup story, a new thing to monitor — for a feature that, for most products, doesn't need vector search on day one.
|
|
24
|
+
- **LangGraph's `store` protocol is a clean seam.** It's designed so storage is swappable — your agent code shouldn't care whether memories live in Postgres, Dynamo, or Redis. This fills the Dynamo gap in that seam.
|
|
25
|
+
- **Memory doesn't have to mean embeddings.** Most products get real value from "the last N things we know about this user," fetched by recency — no vector index required. `memory-layer` starts there, and gives you a clean place to add semantic ranking later if you actually need it.
|
|
26
|
+
|
|
27
|
+
## Features
|
|
28
|
+
|
|
29
|
+
- **`DynamoDBStore`** — a complete `langgraph.store.base.BaseStore` implementation: `get`/`put`/`search` (and their async counterparts), real pagination via DynamoDB's `LastEvaluatedKey` (not a "hope the first page has enough" heuristic), per-item filtering, and per-type TTL (e.g. auto-expire `episodic` memories after 90 days while `semantic` ones never expire).
|
|
30
|
+
- **`SimpleRetrieval`** — fetch a user's N most recent memories and turn them into a ready-to-inject prompt block. No embeddings, no extra dependencies.
|
|
31
|
+
- **`MemoryWriter`** — an LLM-driven extraction step: hand it a conversation, it classifies and persists the facts worth remembering, via `with_structured_output` (a real schema-validated response, not a hand-rolled JSON parser hoping the model didn't wrap the array in a sentence).
|
|
32
|
+
- **Scopes, not just users.** Namespaces are plain tuples (`("user", user_id)`, `("instance", tenant_id)`) — model per-user memory, per-tenant shared context, or your own scope, however your product actually shapes ownership.
|
|
33
|
+
|
|
34
|
+
## Install
|
|
35
|
+
|
|
36
|
+
```bash
|
|
37
|
+
pip install langgraph-dynamodb-store
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
Want LLM-driven extraction? `MemoryWriter` takes any LangChain `BaseChatModel` — bring the one you already use, no extra install needed. `numpy`/`langchain-openai` are only required for the semantic-retrieval extra (see [Roadmap](#roadmap)):
|
|
41
|
+
|
|
42
|
+
```bash
|
|
43
|
+
pip install "langgraph-dynamodb-store[semantic]"
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
## DynamoDB table
|
|
47
|
+
|
|
48
|
+
One table, one GSI. Create it however you provision infra (CDK/Terraform/console) — here's the raw shape via the AWS CLI, if you just want to try it out:
|
|
49
|
+
|
|
50
|
+
```bash
|
|
51
|
+
aws dynamodb create-table \
|
|
52
|
+
--table-name my-app-memories \
|
|
53
|
+
--attribute-definitions \
|
|
54
|
+
AttributeName=owner_id,AttributeType=S \
|
|
55
|
+
AttributeName=memory_id,AttributeType=S \
|
|
56
|
+
AttributeName=created_at,AttributeType=S \
|
|
57
|
+
--key-schema \
|
|
58
|
+
AttributeName=owner_id,KeyType=HASH \
|
|
59
|
+
AttributeName=memory_id,KeyType=RANGE \
|
|
60
|
+
--global-secondary-indexes '[{
|
|
61
|
+
"IndexName": "owner_id-created_at-index",
|
|
62
|
+
"KeySchema": [
|
|
63
|
+
{"AttributeName": "owner_id", "KeyType": "HASH"},
|
|
64
|
+
{"AttributeName": "created_at", "KeyType": "RANGE"}
|
|
65
|
+
],
|
|
66
|
+
"Projection": {"ProjectionType": "ALL"}
|
|
67
|
+
}]' \
|
|
68
|
+
--billing-mode PAY_PER_REQUEST
|
|
69
|
+
|
|
70
|
+
# Optional but recommended — lets episodic memories actually expire instead of
|
|
71
|
+
# accumulating forever. memory-layer sets the `ttl` attribute; DynamoDB does the rest.
|
|
72
|
+
aws dynamodb update-time-to-live \
|
|
73
|
+
--table-name my-app-memories \
|
|
74
|
+
--time-to-live-specification "Enabled=true, AttributeName=ttl"
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
Set `MEMORY_TABLE=my-app-memories` or pass `table_name` explicitly — `DynamoDBStore(table_name="my-app-memories")`.
|
|
78
|
+
|
|
79
|
+
## Quickstart
|
|
80
|
+
|
|
81
|
+
```python
|
|
82
|
+
from memory_layer import DynamoDBStore
|
|
83
|
+
|
|
84
|
+
store = DynamoDBStore(table_name="my-app-memories")
|
|
85
|
+
|
|
86
|
+
namespace = ("user", "user-123")
|
|
87
|
+
|
|
88
|
+
store.put(namespace, "mem-1", {"content": "Prefers responses in Spanish", "type": "semantic"})
|
|
89
|
+
store.put(namespace, "mem-2", {"content": "Reviewed Q3 numbers on 2026-08-01", "type": "episodic"})
|
|
90
|
+
|
|
91
|
+
memories = store.search(namespace, limit=10)
|
|
92
|
+
for m in memories:
|
|
93
|
+
print(m.value["content"])
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
### Inside a LangGraph node
|
|
97
|
+
|
|
98
|
+
LangGraph injects `store` into any node whose signature asks for it:
|
|
99
|
+
|
|
100
|
+
```python
|
|
101
|
+
from langgraph.store.base import BaseStore
|
|
102
|
+
from langchain_core.runnables import RunnableConfig
|
|
103
|
+
|
|
104
|
+
from memory_layer import DynamoDBStore
|
|
105
|
+
from memory_layer.retrieval import SimpleRetrieval
|
|
106
|
+
|
|
107
|
+
store = DynamoDBStore(table_name="my-app-memories")
|
|
108
|
+
graph = builder.compile(checkpointer=checkpointer, store=store)
|
|
109
|
+
|
|
110
|
+
def supervisor_node(state: AgentState, config: RunnableConfig, store: BaseStore) -> dict:
|
|
111
|
+
user_id = state["context"]["user_id"]
|
|
112
|
+
retrieval = SimpleRetrieval(store, limit=5)
|
|
113
|
+
memories = retrieval.fetch(("user", user_id))
|
|
114
|
+
memory_block = retrieval.to_prompt_block(memories)
|
|
115
|
+
# ...inject memory_block into the system prompt
|
|
116
|
+
return {}
|
|
117
|
+
```
|
|
118
|
+
|
|
119
|
+
### Writing memories back
|
|
120
|
+
|
|
121
|
+
```python
|
|
122
|
+
from memory_layer.writer import MemoryWriter
|
|
123
|
+
|
|
124
|
+
writer = MemoryWriter(llm=your_chat_model, store=store)
|
|
125
|
+
|
|
126
|
+
async def memory_writer_node(state: AgentState) -> dict:
|
|
127
|
+
await writer.extract_and_save(
|
|
128
|
+
namespace=("user", state["context"]["user_id"]),
|
|
129
|
+
messages=state["messages"],
|
|
130
|
+
session_id=state["context"]["session_id"],
|
|
131
|
+
)
|
|
132
|
+
return {}
|
|
133
|
+
```
|
|
134
|
+
|
|
135
|
+
The extraction LLM only needs `with_structured_output` support — every major provider's LangChain integration has it.
|
|
136
|
+
|
|
137
|
+
## Memory types
|
|
138
|
+
|
|
139
|
+
| Type | What it's for | Default TTL |
|
|
140
|
+
|---|---|---|
|
|
141
|
+
| `semantic` | Stable preferences and facts ("prefers Spanish", "works on the Acme account") | none |
|
|
142
|
+
| `episodic` | Specific past events ("reviewed Q3 numbers on 2026-08-01") | 90 days |
|
|
143
|
+
| `procedural` | Recurring work patterns ("always starts with a channel breakdown") | none |
|
|
144
|
+
|
|
145
|
+
TTL policy per type lives in `memory_layer.store._TTL_SECONDS_BY_TYPE` — override it if 90 days isn't the right default for your product.
|
|
146
|
+
|
|
147
|
+
## Scopes
|
|
148
|
+
|
|
149
|
+
A namespace is just a tuple — `memory-layer` doesn't prescribe what it means, but the common shapes are:
|
|
150
|
+
|
|
151
|
+
```python
|
|
152
|
+
("user", cognito_sub) # private to one user
|
|
153
|
+
("instance", tenant_id) # shared across every user of one tenant
|
|
154
|
+
```
|
|
155
|
+
|
|
156
|
+
## Roadmap
|
|
157
|
+
|
|
158
|
+
- **Semantic retrieval** — embed memories + query, rank by cosine similarity, for products that outgrow "most recent N" (roughly ~20+ memories per user is where this starts to matter). Lives behind the `semantic` extra so the core library stays dependency-light.
|
|
159
|
+
- **Deduplication on write** — skip persisting a fact that's a near-duplicate of one already stored.
|
|
160
|
+
- **Bring-your-own embeddings backend** — DynamoDB-native cosine similarity to start; pluggable enough to swap in pgvector/a real vector store later if volume ever justifies it.
|
|
161
|
+
|
|
162
|
+
## Development
|
|
163
|
+
|
|
164
|
+
```bash
|
|
165
|
+
pip install -e ".[dev]"
|
|
166
|
+
pytest
|
|
167
|
+
```
|
|
168
|
+
|
|
169
|
+
Tests run against [`moto`](https://github.com/getmoto/moto) — no real AWS account or network access required.
|
|
170
|
+
|
|
171
|
+
## License
|
|
172
|
+
|
|
173
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "langgraph-dynamodb-store"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "A DynamoDB-backed BaseStore for LangGraph — cross-thread memory for agents, no Postgres/pgvector required"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "MIT"
|
|
11
|
+
requires-python = ">=3.11"
|
|
12
|
+
authors = [
|
|
13
|
+
{ name = "Mauricio Neira" },
|
|
14
|
+
]
|
|
15
|
+
keywords = ["langgraph", "langchain", "dynamodb", "memory", "agents", "llm"]
|
|
16
|
+
classifiers = [
|
|
17
|
+
"Development Status :: 3 - Alpha",
|
|
18
|
+
"Intended Audience :: Developers",
|
|
19
|
+
"License :: OSI Approved :: MIT License",
|
|
20
|
+
"Programming Language :: Python :: 3",
|
|
21
|
+
"Programming Language :: Python :: 3.11",
|
|
22
|
+
"Programming Language :: Python :: 3.12",
|
|
23
|
+
"Topic :: Software Development :: Libraries :: Python Modules",
|
|
24
|
+
]
|
|
25
|
+
dependencies = [
|
|
26
|
+
"langgraph>=0.2",
|
|
27
|
+
"langchain-core>=0.3",
|
|
28
|
+
"pydantic>=2.0",
|
|
29
|
+
"boto3>=1.35",
|
|
30
|
+
]
|
|
31
|
+
|
|
32
|
+
[project.optional-dependencies]
|
|
33
|
+
semantic = [
|
|
34
|
+
"langchain-openai>=0.2",
|
|
35
|
+
"numpy>=1.26",
|
|
36
|
+
]
|
|
37
|
+
dev = [
|
|
38
|
+
"pytest>=8.0",
|
|
39
|
+
"pytest-asyncio>=0.23",
|
|
40
|
+
"moto[dynamodb]>=5.0",
|
|
41
|
+
]
|
|
42
|
+
|
|
43
|
+
[project.urls]
|
|
44
|
+
Homepage = "https://github.com/mauriciosneira/memory-layer"
|
|
45
|
+
Issues = "https://github.com/mauriciosneira/memory-layer/issues"
|
|
46
|
+
|
|
47
|
+
# Distribution name (langgraph-dynamodb-store) differs from the importable module
|
|
48
|
+
# (memory_layer) — hatchling's src-layout auto-detection derives the package dir from
|
|
49
|
+
# the project name by default, which would look for src/langgraph_dynamodb_store and
|
|
50
|
+
# fail to find anything. Point it at the real package explicitly instead.
|
|
51
|
+
[tool.hatch.build.targets.wheel]
|
|
52
|
+
packages = ["src/memory_layer"]
|
|
53
|
+
|
|
54
|
+
[tool.pytest.ini_options]
|
|
55
|
+
asyncio_mode = "auto"
|
|
56
|
+
testpaths = ["tests"]
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from langgraph.store.base import BaseStore, SearchItem
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
class SimpleRetrieval:
|
|
7
|
+
"""Fetches N most recent memories by recency. No embeddings required."""
|
|
8
|
+
|
|
9
|
+
def __init__(self, store: BaseStore, limit: int = 10) -> None:
|
|
10
|
+
self._store = store
|
|
11
|
+
self._limit = limit
|
|
12
|
+
|
|
13
|
+
def fetch(
|
|
14
|
+
self,
|
|
15
|
+
namespace: tuple[str, ...],
|
|
16
|
+
*,
|
|
17
|
+
memory_type: str | None = None,
|
|
18
|
+
) -> list[SearchItem]:
|
|
19
|
+
filter_ = {"type": memory_type} if memory_type else None
|
|
20
|
+
return self._store.search(namespace, filter=filter_, limit=self._limit)
|
|
21
|
+
|
|
22
|
+
async def afetch(
|
|
23
|
+
self,
|
|
24
|
+
namespace: tuple[str, ...],
|
|
25
|
+
*,
|
|
26
|
+
memory_type: str | None = None,
|
|
27
|
+
) -> list[SearchItem]:
|
|
28
|
+
filter_ = {"type": memory_type} if memory_type else None
|
|
29
|
+
return await self._store.asearch(namespace, filter=filter_, limit=self._limit)
|
|
30
|
+
|
|
31
|
+
def to_prompt_block(self, memories: list[SearchItem]) -> str:
|
|
32
|
+
if not memories:
|
|
33
|
+
return ""
|
|
34
|
+
lines = "\n".join(f"- {m.value.get('content', '')}" for m in memories)
|
|
35
|
+
return f"## User memory\n{lines}"
|
|
@@ -0,0 +1,159 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
import json
|
|
5
|
+
import os
|
|
6
|
+
from datetime import datetime, timezone
|
|
7
|
+
from typing import Any, Iterable
|
|
8
|
+
|
|
9
|
+
import boto3
|
|
10
|
+
from boto3.dynamodb.conditions import Key
|
|
11
|
+
from langgraph.store.base import (
|
|
12
|
+
BaseStore,
|
|
13
|
+
GetOp,
|
|
14
|
+
Item,
|
|
15
|
+
ListNamespacesOp,
|
|
16
|
+
Op,
|
|
17
|
+
PutOp,
|
|
18
|
+
Result,
|
|
19
|
+
SearchItem,
|
|
20
|
+
SearchOp,
|
|
21
|
+
)
|
|
22
|
+
|
|
23
|
+
_TABLE_NAME = os.environ.get("MEMORY_TABLE", "memory-layer-local-memories")
|
|
24
|
+
_MAX_SEARCH_PAGES = 10
|
|
25
|
+
_TTL_SECONDS_BY_TYPE = {"episodic": 90 * 24 * 60 * 60} # semantic/procedural: no ttl attribute, never expire
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def _owner_id(namespace: tuple[str, ...]) -> str:
|
|
29
|
+
return ":".join(namespace)
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def _parse_namespace(owner_id: str) -> tuple[str, ...]:
|
|
33
|
+
return tuple(owner_id.split(":"))
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def _row_to_item(row: dict[str, Any]) -> Item:
|
|
37
|
+
return Item(
|
|
38
|
+
namespace=_parse_namespace(row["owner_id"]),
|
|
39
|
+
key=row["memory_id"],
|
|
40
|
+
value=json.loads(row["value"]),
|
|
41
|
+
created_at=datetime.fromisoformat(row["created_at"]),
|
|
42
|
+
updated_at=datetime.fromisoformat(row["updated_at"]),
|
|
43
|
+
)
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def _row_to_search_item(row: dict[str, Any]) -> SearchItem:
|
|
47
|
+
return SearchItem(
|
|
48
|
+
namespace=_parse_namespace(row["owner_id"]),
|
|
49
|
+
key=row["memory_id"],
|
|
50
|
+
value=json.loads(row["value"]),
|
|
51
|
+
created_at=datetime.fromisoformat(row["created_at"]),
|
|
52
|
+
updated_at=datetime.fromisoformat(row["updated_at"]),
|
|
53
|
+
score=None,
|
|
54
|
+
)
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def _matches_filter(value: dict[str, Any], filter: dict[str, Any]) -> bool:
|
|
58
|
+
for field, condition in filter.items():
|
|
59
|
+
field_val = value.get(field)
|
|
60
|
+
if isinstance(condition, dict):
|
|
61
|
+
for op, operand in condition.items():
|
|
62
|
+
if op == "$eq" and field_val != operand:
|
|
63
|
+
return False
|
|
64
|
+
elif op == "$ne" and field_val == operand:
|
|
65
|
+
return False
|
|
66
|
+
elif op == "$gt" and not (field_val is not None and field_val > operand):
|
|
67
|
+
return False
|
|
68
|
+
elif op == "$gte" and not (field_val is not None and field_val >= operand):
|
|
69
|
+
return False
|
|
70
|
+
elif op == "$lt" and not (field_val is not None and field_val < operand):
|
|
71
|
+
return False
|
|
72
|
+
elif op == "$lte" and not (field_val is not None and field_val <= operand):
|
|
73
|
+
return False
|
|
74
|
+
elif field_val != condition:
|
|
75
|
+
return False
|
|
76
|
+
return True
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
class DynamoDBStore(BaseStore):
|
|
80
|
+
def __init__(self, table_name: str = _TABLE_NAME) -> None:
|
|
81
|
+
self._table = boto3.resource("dynamodb").Table(table_name)
|
|
82
|
+
|
|
83
|
+
def batch(self, ops: Iterable[Op]) -> list[Result]:
|
|
84
|
+
results: list[Result] = []
|
|
85
|
+
for op in ops:
|
|
86
|
+
if isinstance(op, GetOp):
|
|
87
|
+
results.append(self._handle_get(op))
|
|
88
|
+
elif isinstance(op, PutOp):
|
|
89
|
+
results.append(self._handle_put(op))
|
|
90
|
+
elif isinstance(op, SearchOp):
|
|
91
|
+
results.append(self._handle_search(op))
|
|
92
|
+
elif isinstance(op, ListNamespacesOp):
|
|
93
|
+
results.append([])
|
|
94
|
+
else:
|
|
95
|
+
raise TypeError(f"Unsupported op type: {type(op).__name__}")
|
|
96
|
+
return results
|
|
97
|
+
|
|
98
|
+
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
|
|
99
|
+
return await asyncio.to_thread(self.batch, list(ops))
|
|
100
|
+
|
|
101
|
+
def _handle_get(self, op: GetOp) -> Item | None:
|
|
102
|
+
resp = self._table.get_item(
|
|
103
|
+
Key={"owner_id": _owner_id(op.namespace), "memory_id": op.key}
|
|
104
|
+
)
|
|
105
|
+
row = resp.get("Item")
|
|
106
|
+
return _row_to_item(row) if row else None
|
|
107
|
+
|
|
108
|
+
def _handle_put(self, op: PutOp) -> None:
|
|
109
|
+
pk = {"owner_id": _owner_id(op.namespace), "memory_id": op.key}
|
|
110
|
+
if op.value is None:
|
|
111
|
+
self._table.delete_item(Key=pk)
|
|
112
|
+
return
|
|
113
|
+
|
|
114
|
+
now = datetime.now(timezone.utc).isoformat()
|
|
115
|
+
update_expression = (
|
|
116
|
+
"SET #v = :v, updated_at = :now, "
|
|
117
|
+
"created_at = if_not_exists(created_at, :now)"
|
|
118
|
+
)
|
|
119
|
+
expression_values: dict[str, Any] = {":v": json.dumps(op.value), ":now": now}
|
|
120
|
+
|
|
121
|
+
ttl_seconds = _TTL_SECONDS_BY_TYPE.get(op.value.get("type"))
|
|
122
|
+
if ttl_seconds is not None:
|
|
123
|
+
update_expression += ", #ttl = :ttl"
|
|
124
|
+
expression_values[":ttl"] = int(datetime.now(timezone.utc).timestamp()) + ttl_seconds
|
|
125
|
+
|
|
126
|
+
self._table.update_item(
|
|
127
|
+
Key=pk,
|
|
128
|
+
UpdateExpression=update_expression,
|
|
129
|
+
ExpressionAttributeNames={"#v": "value", "#ttl": "ttl"} if ttl_seconds is not None else {"#v": "value"},
|
|
130
|
+
ExpressionAttributeValues=expression_values,
|
|
131
|
+
)
|
|
132
|
+
|
|
133
|
+
def _handle_search(self, op: SearchOp) -> list[SearchItem]:
|
|
134
|
+
matches: list[dict[str, Any]] = []
|
|
135
|
+
last_evaluated_key: dict[str, Any] | None = None
|
|
136
|
+
pages = 0
|
|
137
|
+
|
|
138
|
+
while len(matches) < op.offset + op.limit and pages < _MAX_SEARCH_PAGES:
|
|
139
|
+
query_kwargs: dict[str, Any] = {
|
|
140
|
+
"IndexName": "owner_id-created_at-index",
|
|
141
|
+
"KeyConditionExpression": Key("owner_id").eq(_owner_id(op.namespace_prefix)),
|
|
142
|
+
"ScanIndexForward": False,
|
|
143
|
+
}
|
|
144
|
+
if last_evaluated_key is not None:
|
|
145
|
+
query_kwargs["ExclusiveStartKey"] = last_evaluated_key
|
|
146
|
+
|
|
147
|
+
resp = self._table.query(**query_kwargs)
|
|
148
|
+
rows = resp.get("Items", [])
|
|
149
|
+
pages += 1
|
|
150
|
+
|
|
151
|
+
if op.filter:
|
|
152
|
+
rows = [r for r in rows if _matches_filter(json.loads(r["value"]), op.filter)]
|
|
153
|
+
matches.extend(rows)
|
|
154
|
+
|
|
155
|
+
last_evaluated_key = resp.get("LastEvaluatedKey")
|
|
156
|
+
if last_evaluated_key is None:
|
|
157
|
+
break
|
|
158
|
+
|
|
159
|
+
return [_row_to_search_item(r) for r in matches[op.offset : op.offset + op.limit]]
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
from typing import Literal, TypedDict, Optional
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
MemoryType = Literal["semantic", "episodic", "procedural"]
|
|
6
|
+
MemoryScope = Literal["user", "instance"]
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class MemoryRecord(TypedDict, total=False):
|
|
10
|
+
content: str
|
|
11
|
+
type: MemoryType
|
|
12
|
+
created_at: str
|
|
13
|
+
updated_at: str
|
|
14
|
+
session_id: str
|
|
15
|
+
score: Optional[float]
|
|
@@ -0,0 +1,75 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from uuid import uuid4
|
|
4
|
+
|
|
5
|
+
from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage
|
|
6
|
+
from langchain_core.language_models import BaseChatModel
|
|
7
|
+
from langgraph.store.base import BaseStore
|
|
8
|
+
from pydantic import BaseModel
|
|
9
|
+
|
|
10
|
+
from .types import MemoryType
|
|
11
|
+
|
|
12
|
+
_EXTRACTION_PROMPT = """You are a memory extraction assistant. Given a conversation, extract facts worth remembering about the user for future sessions.
|
|
13
|
+
|
|
14
|
+
Rules:
|
|
15
|
+
- Only extract facts that would be useful in a *future* conversation — preferences, recurring patterns, key domain context.
|
|
16
|
+
- Ignore facts that are specific to this single request and won't generalise.
|
|
17
|
+
- Each fact must be a short, standalone sentence.
|
|
18
|
+
- Classify each fact as one of: semantic (preferences/facts about the user), episodic (specific past event), procedural (recurring work pattern).
|
|
19
|
+
- If nothing is worth remembering, return no facts."""
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class _ExtractedFact(BaseModel):
|
|
23
|
+
content: str
|
|
24
|
+
type: MemoryType
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
class _ExtractedFacts(BaseModel):
|
|
28
|
+
facts: list[_ExtractedFact]
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
class MemoryWriter:
|
|
32
|
+
def __init__(self, llm: BaseChatModel, store: BaseStore) -> None:
|
|
33
|
+
self._llm = llm
|
|
34
|
+
self._store = store
|
|
35
|
+
|
|
36
|
+
async def extract_and_save(
|
|
37
|
+
self,
|
|
38
|
+
namespace: tuple[str, ...],
|
|
39
|
+
messages: list[BaseMessage],
|
|
40
|
+
session_id: str,
|
|
41
|
+
) -> int:
|
|
42
|
+
conversation = _format_conversation(messages)
|
|
43
|
+
if not conversation:
|
|
44
|
+
return 0
|
|
45
|
+
|
|
46
|
+
extraction_messages = [
|
|
47
|
+
SystemMessage(content=_EXTRACTION_PROMPT),
|
|
48
|
+
HumanMessage(content=conversation),
|
|
49
|
+
]
|
|
50
|
+
# with_structured_output validates the shape via the provider's own tool-calling —
|
|
51
|
+
# no hand-rolled JSON parsing, no risk of a stray sentence around the array breaking it.
|
|
52
|
+
structured_llm = self._llm.with_structured_output(_ExtractedFacts)
|
|
53
|
+
result: _ExtractedFacts = await structured_llm.ainvoke(extraction_messages)
|
|
54
|
+
|
|
55
|
+
for fact in result.facts:
|
|
56
|
+
await self._store.aput(
|
|
57
|
+
namespace,
|
|
58
|
+
str(uuid4()),
|
|
59
|
+
{
|
|
60
|
+
"content": fact.content,
|
|
61
|
+
"type": fact.type,
|
|
62
|
+
"session_id": session_id,
|
|
63
|
+
},
|
|
64
|
+
)
|
|
65
|
+
|
|
66
|
+
return len(result.facts)
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def _format_conversation(messages: list[BaseMessage]) -> str:
|
|
70
|
+
lines: list[str] = []
|
|
71
|
+
for msg in messages:
|
|
72
|
+
role = getattr(msg, "type", "unknown")
|
|
73
|
+
if role in ("human", "ai") and msg.content:
|
|
74
|
+
lines.append(f"{role.upper()}: {msg.content}")
|
|
75
|
+
return "\n".join(lines)
|
|
@@ -0,0 +1,41 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import pytest
|
|
5
|
+
import boto3
|
|
6
|
+
from moto import mock_aws
|
|
7
|
+
|
|
8
|
+
os.environ["MEMORY_TABLE"] = "test-memories"
|
|
9
|
+
os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
|
|
10
|
+
os.environ["AWS_ACCESS_KEY_ID"] = "test"
|
|
11
|
+
os.environ["AWS_SECRET_ACCESS_KEY"] = "test"
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
@pytest.fixture
|
|
15
|
+
def dynamodb_table():
|
|
16
|
+
with mock_aws():
|
|
17
|
+
client = boto3.client("dynamodb", region_name="us-east-1")
|
|
18
|
+
client.create_table(
|
|
19
|
+
TableName="test-memories",
|
|
20
|
+
KeySchema=[
|
|
21
|
+
{"AttributeName": "owner_id", "KeyType": "HASH"},
|
|
22
|
+
{"AttributeName": "memory_id", "KeyType": "RANGE"},
|
|
23
|
+
],
|
|
24
|
+
AttributeDefinitions=[
|
|
25
|
+
{"AttributeName": "owner_id", "AttributeType": "S"},
|
|
26
|
+
{"AttributeName": "memory_id", "AttributeType": "S"},
|
|
27
|
+
{"AttributeName": "created_at", "AttributeType": "S"},
|
|
28
|
+
],
|
|
29
|
+
GlobalSecondaryIndexes=[
|
|
30
|
+
{
|
|
31
|
+
"IndexName": "owner_id-created_at-index",
|
|
32
|
+
"KeySchema": [
|
|
33
|
+
{"AttributeName": "owner_id", "KeyType": "HASH"},
|
|
34
|
+
{"AttributeName": "created_at", "KeyType": "RANGE"},
|
|
35
|
+
],
|
|
36
|
+
"Projection": {"ProjectionType": "ALL"},
|
|
37
|
+
}
|
|
38
|
+
],
|
|
39
|
+
BillingMode="PAY_PER_REQUEST",
|
|
40
|
+
)
|
|
41
|
+
yield boto3.resource("dynamodb", region_name="us-east-1").Table("test-memories")
|
|
@@ -0,0 +1,128 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import time
|
|
4
|
+
|
|
5
|
+
import pytest
|
|
6
|
+
|
|
7
|
+
from memory_layer.store import DynamoDBStore
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
@pytest.fixture
|
|
11
|
+
def store(dynamodb_table):
|
|
12
|
+
return DynamoDBStore(table_name="test-memories")
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
NAMESPACE = ("user", "test-user-123")
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def test_put_and_get(store):
|
|
19
|
+
store.put(NAMESPACE, "mem-1", {"content": "User prefers Spanish", "type": "semantic"})
|
|
20
|
+
item = store.get(NAMESPACE, "mem-1")
|
|
21
|
+
|
|
22
|
+
assert item is not None
|
|
23
|
+
assert item.key == "mem-1"
|
|
24
|
+
assert item.namespace == NAMESPACE
|
|
25
|
+
assert item.value["content"] == "User prefers Spanish"
|
|
26
|
+
assert item.created_at is not None
|
|
27
|
+
assert item.updated_at is not None
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def test_put_preserves_created_at_on_update(store):
|
|
31
|
+
store.put(NAMESPACE, "mem-1", {"content": "original", "type": "semantic"})
|
|
32
|
+
item_before = store.get(NAMESPACE, "mem-1")
|
|
33
|
+
|
|
34
|
+
store.put(NAMESPACE, "mem-1", {"content": "updated", "type": "semantic"})
|
|
35
|
+
item_after = store.get(NAMESPACE, "mem-1")
|
|
36
|
+
|
|
37
|
+
assert item_after.created_at == item_before.created_at
|
|
38
|
+
assert item_after.value["content"] == "updated"
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def test_put_with_none_deletes_item(store):
|
|
42
|
+
store.put(NAMESPACE, "mem-1", {"content": "to be deleted", "type": "semantic"})
|
|
43
|
+
store.put(NAMESPACE, "mem-1", None)
|
|
44
|
+
|
|
45
|
+
assert store.get(NAMESPACE, "mem-1") is None
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def test_get_missing_returns_none(store):
|
|
49
|
+
assert store.get(NAMESPACE, "nonexistent") is None
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def test_search_returns_items(store):
|
|
53
|
+
store.put(NAMESPACE, "mem-1", {"content": "first", "type": "semantic"})
|
|
54
|
+
store.put(NAMESPACE, "mem-2", {"content": "second", "type": "episodic"})
|
|
55
|
+
store.put(NAMESPACE, "mem-3", {"content": "third", "type": "semantic"})
|
|
56
|
+
|
|
57
|
+
results = store.search(NAMESPACE, limit=10)
|
|
58
|
+
|
|
59
|
+
assert len(results) == 3
|
|
60
|
+
keys = {r.key for r in results}
|
|
61
|
+
assert {"mem-1", "mem-2", "mem-3"} == keys
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def test_search_with_type_filter(store):
|
|
65
|
+
store.put(NAMESPACE, "mem-1", {"content": "preference", "type": "semantic"})
|
|
66
|
+
store.put(NAMESPACE, "mem-2", {"content": "past event", "type": "episodic"})
|
|
67
|
+
|
|
68
|
+
results = store.search(NAMESPACE, filter={"type": "semantic"}, limit=10)
|
|
69
|
+
|
|
70
|
+
assert len(results) == 1
|
|
71
|
+
assert results[0].value["type"] == "semantic"
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
def test_search_limit_and_offset(store):
|
|
75
|
+
for i in range(5):
|
|
76
|
+
store.put(NAMESPACE, f"mem-{i}", {"content": f"memory {i}", "type": "semantic"})
|
|
77
|
+
|
|
78
|
+
page1 = store.search(NAMESPACE, limit=2, offset=0)
|
|
79
|
+
page2 = store.search(NAMESPACE, limit=2, offset=2)
|
|
80
|
+
|
|
81
|
+
assert len(page1) == 2
|
|
82
|
+
assert len(page2) == 2
|
|
83
|
+
assert {r.key for r in page1}.isdisjoint({r.key for r in page2})
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def test_search_isolated_by_namespace(store):
|
|
87
|
+
other_namespace = ("user", "other-user")
|
|
88
|
+
store.put(NAMESPACE, "mem-1", {"content": "my memory", "type": "semantic"})
|
|
89
|
+
store.put(other_namespace, "mem-2", {"content": "their memory", "type": "semantic"})
|
|
90
|
+
|
|
91
|
+
results = store.search(NAMESPACE, limit=10)
|
|
92
|
+
|
|
93
|
+
assert len(results) == 1
|
|
94
|
+
assert results[0].key == "mem-1"
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def test_unsupported_op_raises_type_error(store):
|
|
98
|
+
class UnknownOp:
|
|
99
|
+
pass
|
|
100
|
+
|
|
101
|
+
with pytest.raises(TypeError, match="Unsupported op type"):
|
|
102
|
+
store.batch([UnknownOp()])
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
def test_put_episodic_sets_a_ttl(store):
|
|
106
|
+
store.put(NAMESPACE, "mem-1", {"content": "went to the store", "type": "episodic"})
|
|
107
|
+
|
|
108
|
+
raw = store._table.get_item(Key={"owner_id": ":".join(NAMESPACE), "memory_id": "mem-1"})["Item"]
|
|
109
|
+
assert "ttl" in raw
|
|
110
|
+
assert int(raw["ttl"]) > int(time.time())
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
def test_put_semantic_sets_no_ttl(store):
|
|
114
|
+
store.put(NAMESPACE, "mem-1", {"content": "prefers Spanish", "type": "semantic"})
|
|
115
|
+
|
|
116
|
+
raw = store._table.get_item(Key={"owner_id": ":".join(NAMESPACE), "memory_id": "mem-1"})["Item"]
|
|
117
|
+
assert "ttl" not in raw
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def test_search_paginates_past_the_first_page_when_filter_thins_it_out(store):
|
|
121
|
+
for i in range(15):
|
|
122
|
+
memory_type = "semantic" if i == 14 else "episodic"
|
|
123
|
+
store.put(NAMESPACE, f"mem-{i}", {"content": f"memory {i}", "type": memory_type})
|
|
124
|
+
|
|
125
|
+
results = store.search(NAMESPACE, filter={"type": "semantic"}, limit=1)
|
|
126
|
+
|
|
127
|
+
assert len(results) == 1
|
|
128
|
+
assert results[0].value["type"] == "semantic"
|
|
@@ -0,0 +1,86 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import pytest
|
|
4
|
+
from unittest.mock import AsyncMock, MagicMock
|
|
5
|
+
from langchain_core.messages import HumanMessage, AIMessage
|
|
6
|
+
|
|
7
|
+
from memory_layer.writer import MemoryWriter, _ExtractedFact, _ExtractedFacts, _format_conversation
|
|
8
|
+
from memory_layer.store import DynamoDBStore
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
NAMESPACE = ("user", "test-user-123")
|
|
12
|
+
SESSION_ID = "session-abc"
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def _make_llm(facts: list[_ExtractedFact]) -> MagicMock:
|
|
16
|
+
structured = MagicMock()
|
|
17
|
+
structured.ainvoke = AsyncMock(return_value=_ExtractedFacts(facts=facts))
|
|
18
|
+
llm = MagicMock()
|
|
19
|
+
llm.with_structured_output = MagicMock(return_value=structured)
|
|
20
|
+
return llm
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@pytest.fixture
|
|
24
|
+
def store(dynamodb_table):
|
|
25
|
+
return DynamoDBStore(table_name="test-memories")
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
@pytest.mark.asyncio
|
|
29
|
+
async def test_extract_and_save_writes_facts(store):
|
|
30
|
+
facts = [
|
|
31
|
+
_ExtractedFact(content="User prefers Spanish", type="semantic"),
|
|
32
|
+
_ExtractedFact(content="User works with the Acme account", type="semantic"),
|
|
33
|
+
]
|
|
34
|
+
writer = MemoryWriter(llm=_make_llm(facts), store=store)
|
|
35
|
+
messages = [
|
|
36
|
+
HumanMessage(content="Analyze Acme campaigns"),
|
|
37
|
+
AIMessage(content="Here is the analysis..."),
|
|
38
|
+
]
|
|
39
|
+
|
|
40
|
+
count = await writer.extract_and_save(NAMESPACE, messages, SESSION_ID)
|
|
41
|
+
|
|
42
|
+
assert count == 2
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
@pytest.mark.asyncio
|
|
46
|
+
async def test_extract_and_save_empty_conversation_skips_llm():
|
|
47
|
+
llm = _make_llm([])
|
|
48
|
+
writer = MemoryWriter(llm=llm, store=AsyncMock())
|
|
49
|
+
|
|
50
|
+
count = await writer.extract_and_save(NAMESPACE, [], SESSION_ID)
|
|
51
|
+
|
|
52
|
+
assert count == 0
|
|
53
|
+
llm.with_structured_output.assert_not_called()
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
@pytest.mark.asyncio
|
|
57
|
+
async def test_extract_and_save_empty_facts_writes_nothing():
|
|
58
|
+
store = AsyncMock()
|
|
59
|
+
writer = MemoryWriter(llm=_make_llm([]), store=store)
|
|
60
|
+
messages = [HumanMessage(content="Hello"), AIMessage(content="Hi")]
|
|
61
|
+
|
|
62
|
+
count = await writer.extract_and_save(NAMESPACE, messages, SESSION_ID)
|
|
63
|
+
|
|
64
|
+
assert count == 0
|
|
65
|
+
store.aput.assert_not_called()
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
def test_format_conversation_includes_human_and_ai():
|
|
69
|
+
messages = [HumanMessage(content="Hello"), AIMessage(content="Hi there")]
|
|
70
|
+
result = _format_conversation(messages)
|
|
71
|
+
|
|
72
|
+
assert "HUMAN: Hello" in result
|
|
73
|
+
assert "AI: Hi there" in result
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def test_format_conversation_skips_empty_content():
|
|
77
|
+
messages = [HumanMessage(content=""), AIMessage(content="response")]
|
|
78
|
+
result = _format_conversation(messages)
|
|
79
|
+
|
|
80
|
+
assert "HUMAN" not in result
|
|
81
|
+
assert "AI: response" in result
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def test_extracted_fact_rejects_an_invalid_type():
|
|
85
|
+
with pytest.raises(ValueError):
|
|
86
|
+
_ExtractedFact(content="Bad", type="invalid_type")
|