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.
- langgraph_store_core-0.1.0/.gitignore +9 -0
- langgraph_store_core-0.1.0/PKG-INFO +57 -0
- langgraph_store_core-0.1.0/README.md +41 -0
- langgraph_store_core-0.1.0/pyproject.toml +28 -0
- langgraph_store_core-0.1.0/src/langgraph_store_core/__init__.py +349 -0
- langgraph_store_core-0.1.0/src/langgraph_store_core/testing.py +68 -0
- langgraph_store_core-0.1.0/tests/test_core_semantic.py +118 -0
|
@@ -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
|