langgraph-store-core 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.
@@ -0,0 +1,9 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *.egg-info/
4
+ .pytest_cache/
5
+ .venv/
6
+ dist/
7
+ build/
8
+ uv.lock
9
+ .env
@@ -0,0 +1,57 @@
1
+ Metadata-Version: 2.5
2
+ Name: langgraph-store-core
3
+ Version: 0.1.0
4
+ Summary: Shared core for LangGraph long-term-memory (BaseStore) backends — implement a store with four small primitives
5
+ Project-URL: Homepage, https://github.com/skamalj/langgraph-store
6
+ Project-URL: Repository, https://github.com/skamalj/langgraph-store.git
7
+ Author-email: Kamal <skamalj@gmail.com>
8
+ Keywords: agent-memory,basestore,langgraph,long-term-memory,store
9
+ Classifier: License :: OSI Approved :: MIT License
10
+ Classifier: Operating System :: OS Independent
11
+ Classifier: Programming Language :: Python :: 3
12
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
13
+ Requires-Python: >=3.10
14
+ Requires-Dist: langgraph>=0.2.60
15
+ Description-Content-Type: text/markdown
16
+
17
+ # langgraph-store-core
18
+
19
+ Shared core for building [LangGraph](https://langchain-ai.github.io/langgraph/) long-term-memory stores (`BaseStore`). Implement a backend by subclassing `KVStore` and supplying four small primitives — the core handles `batch`/`abatch`, timestamps, search filters, namespace prefix/suffix matching, `list_namespaces`, and **semantic search**.
20
+
21
+ ```python
22
+ from langgraph_store_core import KVStore
23
+
24
+ class MyStore(KVStore):
25
+ def _read(self, prefix, key): ... # -> row | None
26
+ def _write(self, row): ... # row = {prefix,key,value,created_at,updated_at,embedding}
27
+ def _remove(self, prefix, key): ...
28
+ def _scan(self, prefix): ... # rows whose prefix starts with `prefix`
29
+
30
+ # optional — native vector search; the default ranks `_scan` rows by cosine in Python
31
+ def _vector_search(self, prefix, vector, filter, limit): ... # -> [(row, score), ...]
32
+ ```
33
+
34
+ ## Semantic search
35
+
36
+ Pass LangGraph's `IndexConfig` to any store built on the core and it embeds the configured fields on `put` and answers `search(namespace, query="...")` ranked by cosine similarity, with `SearchItem.score` populated:
37
+
38
+ ```python
39
+ store = MyStore(index={"dims": 1024, "embed": embedder, "fields": ["text"]})
40
+ store.put(("memories", "u1"), "k", {"text": "the user loves sushi"})
41
+ hits = store.search(("memories", "u1"), query="what food does the user like?")
42
+ ```
43
+
44
+ - `embed`: a LangChain `Embeddings`, a `list[str] -> list[list[float]]` callable, or a provider string (`"openai:text-embedding-3-small"`). `bedrock_titan_embeddings()` is included as a ready-made callable (needs `boto3`).
45
+ - `fields`: JSON paths to embed; default `["$"]` (whole value). `put(..., index=False)` skips one item; `put(..., index=["title"])` overrides the fields.
46
+ - A store without `index` is filter-only and ignores `query`, exactly as LangGraph documents.
47
+ - Backends override `_vector_search` for native ANN (pgvector, Cosmos `VectorDistance`, Firestore `find_nearest`, DynamoDB `SearchVectors`); the core re-checks namespace and filter on what comes back, so a backend may push down only part of the filter.
48
+
49
+ `langgraph_store_core.testing` ships `FakeEmbeddings` (deterministic, no network) and `MemoryKVStore` for tests.
50
+
51
+ Concrete backends: [`langgraph-store-dynamodb`](https://pypi.org/project/langgraph-store-dynamodb/), [`langgraph-store-postgres`](https://pypi.org/project/langgraph-store-postgres/), [`langgraph-store-cosmosdb`](https://pypi.org/project/langgraph-store-cosmosdb/), [`langgraph-store-firestore`](https://pypi.org/project/langgraph-store-firestore/).
52
+
53
+ Docs: <https://skamalj.github.io/agentstate-reducer/>
54
+
55
+ ## License
56
+
57
+ MIT
@@ -0,0 +1,41 @@
1
+ # langgraph-store-core
2
+
3
+ Shared core for building [LangGraph](https://langchain-ai.github.io/langgraph/) long-term-memory stores (`BaseStore`). Implement a backend by subclassing `KVStore` and supplying four small primitives — the core handles `batch`/`abatch`, timestamps, search filters, namespace prefix/suffix matching, `list_namespaces`, and **semantic search**.
4
+
5
+ ```python
6
+ from langgraph_store_core import KVStore
7
+
8
+ class MyStore(KVStore):
9
+ def _read(self, prefix, key): ... # -> row | None
10
+ def _write(self, row): ... # row = {prefix,key,value,created_at,updated_at,embedding}
11
+ def _remove(self, prefix, key): ...
12
+ def _scan(self, prefix): ... # rows whose prefix starts with `prefix`
13
+
14
+ # optional — native vector search; the default ranks `_scan` rows by cosine in Python
15
+ def _vector_search(self, prefix, vector, filter, limit): ... # -> [(row, score), ...]
16
+ ```
17
+
18
+ ## Semantic search
19
+
20
+ Pass LangGraph's `IndexConfig` to any store built on the core and it embeds the configured fields on `put` and answers `search(namespace, query="...")` ranked by cosine similarity, with `SearchItem.score` populated:
21
+
22
+ ```python
23
+ store = MyStore(index={"dims": 1024, "embed": embedder, "fields": ["text"]})
24
+ store.put(("memories", "u1"), "k", {"text": "the user loves sushi"})
25
+ hits = store.search(("memories", "u1"), query="what food does the user like?")
26
+ ```
27
+
28
+ - `embed`: a LangChain `Embeddings`, a `list[str] -> list[list[float]]` callable, or a provider string (`"openai:text-embedding-3-small"`). `bedrock_titan_embeddings()` is included as a ready-made callable (needs `boto3`).
29
+ - `fields`: JSON paths to embed; default `["$"]` (whole value). `put(..., index=False)` skips one item; `put(..., index=["title"])` overrides the fields.
30
+ - A store without `index` is filter-only and ignores `query`, exactly as LangGraph documents.
31
+ - Backends override `_vector_search` for native ANN (pgvector, Cosmos `VectorDistance`, Firestore `find_nearest`, DynamoDB `SearchVectors`); the core re-checks namespace and filter on what comes back, so a backend may push down only part of the filter.
32
+
33
+ `langgraph_store_core.testing` ships `FakeEmbeddings` (deterministic, no network) and `MemoryKVStore` for tests.
34
+
35
+ Concrete backends: [`langgraph-store-dynamodb`](https://pypi.org/project/langgraph-store-dynamodb/), [`langgraph-store-postgres`](https://pypi.org/project/langgraph-store-postgres/), [`langgraph-store-cosmosdb`](https://pypi.org/project/langgraph-store-cosmosdb/), [`langgraph-store-firestore`](https://pypi.org/project/langgraph-store-firestore/).
36
+
37
+ Docs: <https://skamalj.github.io/agentstate-reducer/>
38
+
39
+ ## License
40
+
41
+ MIT
@@ -0,0 +1,28 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "langgraph-store-core"
7
+ version = "0.1.0"
8
+ description = "Shared core for LangGraph long-term-memory (BaseStore) backends — implement a store with four small primitives"
9
+ authors = [{name = "Kamal", email = "skamalj@gmail.com"}]
10
+ readme = "README.md"
11
+ requires-python = ">=3.10"
12
+ dependencies = [
13
+ "langgraph>=0.2.60",
14
+ ]
15
+ keywords = ["langgraph", "store", "basestore", "long-term-memory", "agent-memory"]
16
+ classifiers = [
17
+ "Programming Language :: Python :: 3",
18
+ "License :: OSI Approved :: MIT License",
19
+ "Operating System :: OS Independent",
20
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
21
+ ]
22
+
23
+ [project.urls]
24
+ Homepage = "https://github.com/skamalj/langgraph-store"
25
+ Repository = "https://github.com/skamalj/langgraph-store.git"
26
+
27
+ [tool.hatch.build.targets.wheel]
28
+ packages = ["src/langgraph_store_core"]
@@ -0,0 +1,349 @@
1
+ """Shared core for LangGraph long-term-memory ``BaseStore`` backends.
2
+
3
+ ``KVStore`` implements the whole ``BaseStore`` surface — ``batch`` / ``abatch``
4
+ dispatch, timestamp handling, filter evaluation, namespace prefix/suffix matching,
5
+ ``list_namespaces`` and **semantic search** — on top of a handful of primitives
6
+ each backend supplies:
7
+
8
+ _read(prefix, key) -> row | None
9
+ _write(row) -> None
10
+ _remove(prefix, key)-> None
11
+ _scan(prefix) -> list[row]
12
+ _vector_search(prefix, vector, filter, limit) -> list[(row, score)] (optional)
13
+
14
+ A ``row`` is ``{"prefix": str, "key": str, "value": dict, "created_at": iso,
15
+ "updated_at": iso, "embedding": list[float] | None}``. Namespaces (tuples) are
16
+ joined with a unit-separator into the ``prefix`` string.
17
+
18
+ Semantic search
19
+ ---------------
20
+ Construct a store with LangGraph's ``IndexConfig`` (``dims``, ``embed``,
21
+ ``fields``) and it embeds the configured fields of every value on ``put`` and
22
+ answers ``search(namespace, query="...")`` by vector similarity, populating
23
+ ``SearchItem.score``. Backends with native vector search override
24
+ ``_vector_search``; the default implementation scans the prefix and ranks by
25
+ cosine similarity in Python, so every backend gets semantic search even before
26
+ it has a native path. Without an ``IndexConfig`` the store is filter-only and
27
+ ``query`` is ignored, exactly as LangGraph documents.
28
+ """
29
+
30
+ from __future__ import annotations
31
+
32
+ import asyncio
33
+ import json
34
+ import math
35
+ from datetime import datetime, timezone
36
+ from typing import Any, Callable, Iterable, List, Optional, Sequence, Tuple
37
+
38
+ from langgraph.store.base import (
39
+ BaseStore,
40
+ GetOp,
41
+ IndexConfig,
42
+ Item,
43
+ ListNamespacesOp,
44
+ Op,
45
+ PutOp,
46
+ Result,
47
+ SearchItem,
48
+ SearchOp,
49
+ )
50
+ from langgraph.store.base.embed import ensure_embeddings, get_text_at_path, tokenize_path
51
+
52
+ __all__ = [
53
+ "KVStore",
54
+ "IndexConfig",
55
+ "NS_DELIM",
56
+ "ns_to_str",
57
+ "str_to_ns",
58
+ "cosine_similarity",
59
+ "bedrock_titan_embeddings",
60
+ ]
61
+
62
+ NS_DELIM = "\x1f" # unit separator — not expected inside namespace parts
63
+
64
+
65
+ def ns_to_str(namespace: tuple) -> str:
66
+ return NS_DELIM.join(namespace)
67
+
68
+
69
+ def str_to_ns(prefix: str) -> tuple:
70
+ return tuple(prefix.split(NS_DELIM)) if prefix else ()
71
+
72
+
73
+ def _now_iso() -> str:
74
+ return datetime.now(timezone.utc).isoformat()
75
+
76
+
77
+ def _as_dt(value: Any) -> datetime:
78
+ return value if isinstance(value, datetime) else datetime.fromisoformat(value)
79
+
80
+
81
+ def cosine_similarity(a: Sequence[float], b: Sequence[float]) -> float:
82
+ dot = sum(x * y for x, y in zip(a, b))
83
+ na = math.sqrt(sum(x * x for x in a))
84
+ nb = math.sqrt(sum(y * y for y in b))
85
+ if na == 0.0 or nb == 0.0:
86
+ return 0.0
87
+ return dot / (na * nb)
88
+
89
+
90
+ def bedrock_titan_embeddings(
91
+ *,
92
+ model_id: str = "amazon.titan-embed-text-v2:0",
93
+ dimensions: int = 1024,
94
+ normalize: bool = True,
95
+ region_name: Optional[str] = None,
96
+ client: Any = None,
97
+ ) -> Callable[[List[str]], List[List[float]]]:
98
+ """An ``IndexConfig["embed"]`` callable backed by Amazon Bedrock Titan Text v2.
99
+
100
+ Titan v2 supports 256 / 512 / 1024 dimensions; ``dimensions`` must equal the
101
+ store's ``IndexConfig["dims"]``. Requires ``boto3`` (imported lazily).
102
+ """
103
+ import boto3 # noqa: WPS433 — optional dependency
104
+
105
+ runtime = client or boto3.client("bedrock-runtime", region_name=region_name)
106
+
107
+ def embed(texts: List[str]) -> List[List[float]]:
108
+ out: List[List[float]] = []
109
+ for text in texts:
110
+ body = json.dumps({"inputText": text, "dimensions": dimensions, "normalize": normalize})
111
+ resp = runtime.invoke_model(modelId=model_id, body=body)
112
+ out.append(json.loads(resp["body"].read())["embedding"])
113
+ return out
114
+
115
+ return embed
116
+
117
+
118
+ def _matches_filter(value: dict, filter: Optional[dict]) -> bool:
119
+ """Evaluate a LangGraph search ``filter`` against a stored value.
120
+
121
+ Supports plain equality and the common Mongo-style operators
122
+ ($eq/$ne/$gt/$gte/$lt/$lte/$in/$nin).
123
+ """
124
+ if not filter:
125
+ return True
126
+ for field, expected in filter.items():
127
+ actual = value.get(field)
128
+ if isinstance(expected, dict):
129
+ for op, operand in expected.items():
130
+ if op == "$eq" and actual != operand:
131
+ return False
132
+ elif op == "$ne" and actual == operand:
133
+ return False
134
+ elif op == "$gt" and not (actual is not None and actual > operand):
135
+ return False
136
+ elif op == "$gte" and not (actual is not None and actual >= operand):
137
+ return False
138
+ elif op == "$lt" and not (actual is not None and actual < operand):
139
+ return False
140
+ elif op == "$lte" and not (actual is not None and actual <= operand):
141
+ return False
142
+ elif op == "$in" and actual not in operand:
143
+ return False
144
+ elif op == "$nin" and actual in operand:
145
+ return False
146
+ elif op not in ("$eq", "$ne", "$gt", "$gte", "$lt", "$lte", "$in", "$nin"):
147
+ return False
148
+ elif actual != expected:
149
+ return False
150
+ return True
151
+
152
+
153
+ def _match_path(ns_slice: tuple, path: tuple) -> bool:
154
+ if len(ns_slice) != len(path):
155
+ return False
156
+ return all(p == "*" or p == n for n, p in zip(ns_slice, path))
157
+
158
+
159
+ def _match_condition(namespace: tuple, condition: Any) -> bool:
160
+ path = tuple(condition.path)
161
+ if len(path) > len(namespace):
162
+ return False
163
+ if condition.match_type == "suffix":
164
+ return _match_path(namespace[len(namespace) - len(path):], path)
165
+ return _match_path(namespace[: len(path)], path) # "prefix" (default)
166
+
167
+
168
+ def _in_prefix(row_prefix: str, prefix: str) -> bool:
169
+ return prefix == "" or row_prefix == prefix or row_prefix.startswith(prefix + NS_DELIM)
170
+
171
+
172
+ class KVStore(BaseStore):
173
+ """A ``BaseStore`` implemented over a few datastore primitives.
174
+
175
+ Subclasses implement ``_read`` / ``_write`` / ``_remove`` / ``_scan`` and,
176
+ for native semantic search, ``_vector_search``; this class provides
177
+ everything the LangGraph store protocol requires, including the async
178
+ surface (primitives run in a thread executor).
179
+
180
+ Args:
181
+ index: Optional LangGraph ``IndexConfig`` (``dims``, ``embed``,
182
+ ``fields``). When given, values are embedded on ``put`` and
183
+ ``search(query=...)`` ranks by similarity. ``embed`` may be a
184
+ LangChain ``Embeddings``, a ``list[str] -> list[list[float]]``
185
+ callable, or a provider string LangChain can resolve.
186
+ """
187
+
188
+ supports_native_vector_search: bool = False
189
+
190
+ def __init__(self, *, index: Optional[IndexConfig] = None) -> None:
191
+ super().__init__()
192
+ self.index_config: Optional[IndexConfig] = None
193
+ self.embeddings = None
194
+ self._index_paths: List[Tuple[str, list]] = []
195
+ if index is not None:
196
+ self.index_config = dict(index) # type: ignore[assignment]
197
+ self.embeddings = ensure_embeddings(index.get("embed"))
198
+ fields = index.get("fields") or ["$"]
199
+ self._index_paths = [(f, tokenize_path(f)) for f in fields]
200
+
201
+ @property
202
+ def dims(self) -> Optional[int]:
203
+ return int(self.index_config["dims"]) if self.index_config else None
204
+
205
+ # ------------------------------------------------------------------ primitives
206
+ def _read(self, prefix: str, key: str) -> Optional[dict]:
207
+ raise NotImplementedError
208
+
209
+ def _write(self, row: dict) -> None:
210
+ raise NotImplementedError
211
+
212
+ def _remove(self, prefix: str, key: str) -> None:
213
+ raise NotImplementedError
214
+
215
+ def _scan(self, prefix: str) -> List[dict]:
216
+ """Return every row whose ``prefix`` starts with ``prefix`` (all if "")."""
217
+ raise NotImplementedError
218
+
219
+ def _vector_search(
220
+ self, prefix: str, vector: List[float], filter: Optional[dict], limit: int
221
+ ) -> List[Tuple[dict, float]]:
222
+ """Return up to ``limit`` ``(row, score)`` pairs, most similar first.
223
+
224
+ ``score`` is cosine similarity in ``[-1, 1]`` (higher is better).
225
+ Backends with native vector search override this; the default scans the
226
+ prefix and ranks in Python, which is correct for every backend and fine
227
+ for small namespaces.
228
+ """
229
+ scored: List[Tuple[dict, float]] = []
230
+ for row in self._scan(prefix):
231
+ emb = row.get("embedding")
232
+ if not emb or not _in_prefix(row["prefix"], prefix):
233
+ continue
234
+ if not _matches_filter(row["value"], filter):
235
+ continue
236
+ scored.append((row, cosine_similarity(vector, emb)))
237
+ scored.sort(key=lambda rs: -rs[1])
238
+ return scored[:limit]
239
+
240
+ # ------------------------------------------------------------------ embedding
241
+ def _texts_for(self, value: dict, index: Any) -> List[str]:
242
+ """Texts to embed for a value, honouring the per-put ``index`` override."""
243
+ fields = [f for f, _ in self._index_paths] if index is None else list(index)
244
+ texts: List[str] = []
245
+ for field in fields:
246
+ # Pass the raw field string: get_text_at_path tokenizes it itself and
247
+ # handles "$" (whole value as JSON), which a pre-tokenized path does not.
248
+ texts.extend(get_text_at_path(value, field))
249
+ return [t for t in texts if t and t.strip()]
250
+
251
+ def _embed_value(self, value: dict, index: Any) -> Optional[List[float]]:
252
+ if self.embeddings is None or index is False:
253
+ return None
254
+ texts = self._texts_for(value, index)
255
+ if not texts:
256
+ return None
257
+ # One vector per item: embed the concatenation of the indexed fields.
258
+ return list(self.embeddings.embed_documents(["\n".join(texts)])[0])
259
+
260
+ def _embed_query(self, query: str) -> List[float]:
261
+ assert self.embeddings is not None
262
+ return list(self.embeddings.embed_query(query))
263
+
264
+ # ------------------------------------------------------------------ helpers
265
+ def _to_item(self, row: dict, search: bool = False, score: Optional[float] = None):
266
+ namespace = str_to_ns(row["prefix"])
267
+ created, updated = _as_dt(row["created_at"]), _as_dt(row["updated_at"])
268
+ if search:
269
+ return SearchItem(
270
+ namespace=namespace, key=row["key"], value=row["value"],
271
+ created_at=created, updated_at=updated, score=score,
272
+ )
273
+ return Item(
274
+ value=row["value"], key=row["key"], namespace=namespace,
275
+ created_at=created, updated_at=updated,
276
+ )
277
+
278
+ def _op_get(self, namespace: tuple, key: str) -> Optional[Item]:
279
+ row = self._read(ns_to_str(namespace), key)
280
+ return self._to_item(row) if row else None
281
+
282
+ def _op_put(self, namespace: tuple, key: str, value: dict, index: Any = None) -> None:
283
+ prefix = ns_to_str(namespace)
284
+ existing = self._read(prefix, key)
285
+ now = _now_iso()
286
+ self._write({
287
+ "prefix": prefix, "key": key, "value": value,
288
+ "created_at": existing["created_at"] if existing else now,
289
+ "updated_at": now,
290
+ "embedding": self._embed_value(value, index),
291
+ })
292
+
293
+ def _op_delete(self, namespace: tuple, key: str) -> None:
294
+ self._remove(ns_to_str(namespace), key)
295
+
296
+ def _op_search(self, namespace_prefix, filter, limit, offset, query=None) -> List[SearchItem]:
297
+ prefix = ns_to_str(namespace_prefix)
298
+ if query and self.embeddings is not None:
299
+ vector = self._embed_query(query)
300
+ hits = self._vector_search(prefix, vector, filter, (limit or 10) + (offset or 0))
301
+ # Defensive re-check: a backend may apply only part of the filter natively.
302
+ hits = [(r, s) for r, s in hits if _in_prefix(r["prefix"], prefix) and _matches_filter(r["value"], filter)]
303
+ hits = hits[offset : (offset + limit) if limit is not None else None]
304
+ return [self._to_item(r, search=True, score=s) for r, s in hits]
305
+ rows = [r for r in self._scan(prefix) if _in_prefix(r["prefix"], prefix)]
306
+ rows = [r for r in rows if _matches_filter(r["value"], filter)]
307
+ rows.sort(key=lambda r: (r["prefix"], r["key"]))
308
+ rows = rows[offset : (offset + limit) if limit is not None else None]
309
+ return [self._to_item(r, search=True) for r in rows]
310
+
311
+ def _op_list_namespaces(self, match_conditions, max_depth, limit, offset) -> List[tuple]:
312
+ conditions = match_conditions or ()
313
+ seen = set()
314
+ for row in self._scan(""):
315
+ namespace = str_to_ns(row["prefix"])
316
+ if all(_match_condition(namespace, c) for c in conditions):
317
+ seen.add(namespace[:max_depth] if max_depth is not None else namespace)
318
+ ordered = sorted(seen)
319
+ return ordered[offset : (offset + limit) if limit is not None else None]
320
+
321
+ # ------------------------------------------------------------------ BaseStore
322
+ def batch(self, ops: Iterable[Op]) -> List[Result]:
323
+ results: List[Result] = []
324
+ for op in ops:
325
+ if isinstance(op, GetOp):
326
+ results.append(self._op_get(op.namespace, op.key))
327
+ elif isinstance(op, PutOp):
328
+ if op.value is None:
329
+ self._op_delete(op.namespace, op.key)
330
+ else:
331
+ self._op_put(op.namespace, op.key, op.value, getattr(op, "index", None))
332
+ results.append(None)
333
+ elif isinstance(op, SearchOp):
334
+ results.append(
335
+ self._op_search(
336
+ op.namespace_prefix, op.filter, op.limit, op.offset,
337
+ getattr(op, "query", None),
338
+ )
339
+ )
340
+ elif isinstance(op, ListNamespacesOp):
341
+ results.append(
342
+ self._op_list_namespaces(op.match_conditions, op.max_depth, op.limit, op.offset)
343
+ )
344
+ else: # pragma: no cover
345
+ raise NotImplementedError(f"Unsupported op: {type(op)}")
346
+ return results
347
+
348
+ async def abatch(self, ops: Iterable[Op]) -> List[Result]:
349
+ return await asyncio.get_running_loop().run_in_executor(None, self.batch, list(ops))
@@ -0,0 +1,68 @@
1
+ """Test helpers shared by the provider test suites (and useful for users' tests).
2
+
3
+ ``FakeEmbeddings`` is a deterministic, dependency-free embedder: a hashed
4
+ bag-of-words projected into ``dims`` dimensions and L2-normalised. Texts that
5
+ share words are close, texts that do not are far, and no network call is made.
6
+ Good enough to exercise a store's vector path end to end.
7
+
8
+ ``MemoryKVStore`` is the smallest possible ``KVStore`` — a dict — used to test
9
+ the core semantic-search machinery without any backend.
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ import hashlib
15
+ import math
16
+ import re
17
+ from typing import Dict, List, Optional
18
+
19
+ from . import KVStore
20
+
21
+ __all__ = ["FakeEmbeddings", "MemoryKVStore"]
22
+
23
+ _WORD = re.compile(r"[a-z0-9]+")
24
+
25
+
26
+ class FakeEmbeddings:
27
+ """Deterministic hashed bag-of-words embedder with the LangChain interface."""
28
+
29
+ def __init__(self, dims: int = 64) -> None:
30
+ self.dims = dims
31
+
32
+ def _embed(self, text: str) -> List[float]:
33
+ vec = [0.0] * self.dims
34
+ for word in _WORD.findall(text.lower()):
35
+ h = int(hashlib.md5(word.encode()).hexdigest(), 16)
36
+ vec[h % self.dims] += 1.0
37
+ norm = math.sqrt(sum(v * v for v in vec))
38
+ return [v / norm for v in vec] if norm else [1.0 / math.sqrt(self.dims)] * self.dims
39
+
40
+ def embed_documents(self, texts: List[str]) -> List[List[float]]:
41
+ return [self._embed(t) for t in texts]
42
+
43
+ def embed_query(self, text: str) -> List[float]:
44
+ return self._embed(text)
45
+
46
+ def __call__(self, texts: List[str]) -> List[List[float]]: # callable form too
47
+ return self.embed_documents(texts)
48
+
49
+
50
+ class MemoryKVStore(KVStore):
51
+ """In-process ``KVStore`` over a dict; relies on the core's default vector search."""
52
+
53
+ def __init__(self, **kwargs) -> None:
54
+ super().__init__(**kwargs)
55
+ self._rows: Dict[tuple, dict] = {}
56
+
57
+ def _read(self, prefix: str, key: str) -> Optional[dict]:
58
+ row = self._rows.get((prefix, key))
59
+ return dict(row) if row else None
60
+
61
+ def _write(self, row: dict) -> None:
62
+ self._rows[(row["prefix"], row["key"])] = dict(row)
63
+
64
+ def _remove(self, prefix: str, key: str) -> None:
65
+ self._rows.pop((prefix, key), None)
66
+
67
+ def _scan(self, prefix: str) -> List[dict]:
68
+ return [dict(r) for (p, _), r in self._rows.items() if p.startswith(prefix)]
@@ -0,0 +1,118 @@
1
+ """Core semantic-search machinery, exercised on the dict-backed MemoryKVStore.
2
+
3
+ No backend, no network: FakeEmbeddings is deterministic. Everything here holds
4
+ for every provider because it is the shared code path.
5
+ """
6
+ import pytest
7
+
8
+ from langgraph_store_core import KVStore, cosine_similarity
9
+ from langgraph_store_core.testing import FakeEmbeddings, MemoryKVStore
10
+
11
+ DIMS = 32
12
+
13
+
14
+ @pytest.fixture()
15
+ def store():
16
+ return MemoryKVStore(index={"dims": DIMS, "embed": FakeEmbeddings(DIMS), "fields": ["text"]})
17
+
18
+
19
+ def _seed(store):
20
+ store.put(("memories", "kamal"), "food", {"text": "the user loves sushi and japanese food", "kind": "pref"})
21
+ store.put(("memories", "kamal"), "lang", {"text": "the user writes python every day", "kind": "fact"})
22
+ store.put(("memories", "kamal"), "city", {"text": "the user lives by the sea in a big city", "kind": "fact"})
23
+ store.put(("memories", "priya"), "food", {"text": "loves sushi", "kind": "pref"})
24
+
25
+
26
+ def test_put_embeds_configured_fields(store):
27
+ store.put(("n",), "k", {"text": "hello world", "other": "ignored"})
28
+ row = store._read("n", "k")
29
+ assert row["embedding"] is not None and len(row["embedding"]) == DIMS
30
+
31
+
32
+ def test_put_index_false_skips_embedding(store):
33
+ store.put(("n",), "k", {"text": "hello"}, index=False)
34
+ assert store._read("n", "k")["embedding"] is None
35
+
36
+
37
+ def test_put_index_override_fields(store):
38
+ store.put(("n",), "k", {"text": "aaa", "title": "sushi japanese"}, index=["title"])
39
+ hits = store.search(("n",), query="sushi")
40
+ assert hits and hits[0].key == "k" and hits[0].score > 0
41
+
42
+
43
+ def test_query_ranks_by_similarity_and_sets_score(store):
44
+ _seed(store)
45
+ hits = store.search(("memories", "kamal"), query="sushi and japanese food")
46
+ assert hits[0].key == "food"
47
+ assert all(h.score is not None for h in hits)
48
+ assert hits[0].score >= hits[-1].score
49
+
50
+
51
+ def test_query_respects_namespace_prefix(store):
52
+ _seed(store)
53
+ keys = {(h.namespace, h.key) for h in store.search(("memories",), query="sushi")}
54
+ assert (("memories", "kamal"), "food") in keys and (("memories", "priya"), "food") in keys
55
+ only_kamal = store.search(("memories", "kamal"), query="sushi")
56
+ assert all(h.namespace == ("memories", "kamal") for h in only_kamal)
57
+
58
+
59
+ def test_query_combined_with_filter(store):
60
+ _seed(store)
61
+ hits = store.search(("memories", "kamal"), query="user", filter={"kind": "fact"})
62
+ assert hits and all(h.value["kind"] == "fact" for h in hits)
63
+
64
+
65
+ def test_query_limit_and_offset(store):
66
+ _seed(store)
67
+ first = store.search(("memories", "kamal"), query="user", limit=2)
68
+ rest = store.search(("memories", "kamal"), query="user", limit=2, offset=2)
69
+ assert len(first) == 2 and len(rest) == 1
70
+ assert not {h.key for h in first} & {h.key for h in rest}
71
+
72
+
73
+ def test_search_without_query_is_filter_only_no_score(store):
74
+ _seed(store)
75
+ hits = store.search(("memories", "kamal"))
76
+ assert [h.key for h in hits] == ["city", "food", "lang"]
77
+ assert all(h.score is None for h in hits)
78
+
79
+
80
+ def test_store_without_index_ignores_query():
81
+ s = MemoryKVStore()
82
+ s.put(("n",), "a", {"text": "sushi"})
83
+ s.put(("n",), "b", {"text": "python"})
84
+ hits = s.search(("n",), query="sushi")
85
+ assert [h.key for h in hits] == ["a", "b"] and all(h.score is None for h in hits)
86
+ assert s._read("n", "a")["embedding"] is None
87
+
88
+
89
+ def test_default_fields_embed_whole_value():
90
+ s = MemoryKVStore(index={"dims": DIMS, "embed": FakeEmbeddings(DIMS)})
91
+ s.put(("n",), "a", {"anything": "sushi rolls"})
92
+ s.put(("n",), "b", {"anything": "python code"})
93
+ assert s.search(("n",), query="sushi")[0].key == "a"
94
+
95
+
96
+ def test_callable_embed_is_accepted():
97
+ fe = FakeEmbeddings(DIMS)
98
+ s = MemoryKVStore(index={"dims": DIMS, "embed": fe.embed_documents, "fields": ["text"]})
99
+ s.put(("n",), "a", {"text": "sushi"})
100
+ assert s.search(("n",), query="sushi")[0].score > 0.9
101
+
102
+
103
+ def test_update_reembeds(store):
104
+ store.put(("n",), "k", {"text": "sushi"})
105
+ store.put(("n",), "k", {"text": "python"})
106
+ assert store.search(("n",), query="python")[0].score > store.search(("n",), query="sushi")[0].score
107
+
108
+
109
+ async def test_async_query(store):
110
+ _seed(store)
111
+ hits = await store.asearch(("memories", "kamal"), query="sushi")
112
+ assert hits[0].key == "food"
113
+
114
+
115
+ def test_cosine_similarity():
116
+ assert cosine_similarity([1, 0], [1, 0]) == pytest.approx(1.0)
117
+ assert cosine_similarity([1, 0], [0, 1]) == pytest.approx(0.0)
118
+ assert cosine_similarity([0, 0], [1, 1]) == 0.0