langgraph-store-firestore 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_firestore-0.1.0/.gitignore +9 -0
- langgraph_store_firestore-0.1.0/PKG-INFO +72 -0
- langgraph_store_firestore-0.1.0/README.md +55 -0
- langgraph_store_firestore-0.1.0/pyproject.toml +35 -0
- langgraph_store_firestore-0.1.0/src/langgraph_store_firestore/__init__.py +171 -0
- langgraph_store_firestore-0.1.0/tests/test_firestore_semantic.py +118 -0
- langgraph_store_firestore-0.1.0/tests/test_firestore_store.py +85 -0
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: langgraph-store-firestore
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Google Firestore long-term-memory store (BaseStore) for LangGraph — namespaced key/value memory with filters, list_namespaces and native semantic (Firestore vector) search
|
|
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,firestore,gcp,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: google-cloud-firestore
|
|
15
|
+
Requires-Dist: langgraph-store-core>=0.1.0
|
|
16
|
+
Description-Content-Type: text/markdown
|
|
17
|
+
|
|
18
|
+
# langgraph-store-firestore
|
|
19
|
+
|
|
20
|
+
A **Google Firestore** long-term-memory store (`BaseStore`) for [LangGraph](https://langchain-ai.github.io/langgraph/) — namespaced key/value memory with prefix search, filters, `list_namespaces`, and **native semantic search via Firestore vector search**.
|
|
21
|
+
|
|
22
|
+
```bash
|
|
23
|
+
pip install langgraph-store-firestore
|
|
24
|
+
```
|
|
25
|
+
|
|
26
|
+
```python
|
|
27
|
+
from langgraph_store_firestore import FirestoreStore
|
|
28
|
+
|
|
29
|
+
store = FirestoreStore(project_id="my-gcp-project")
|
|
30
|
+
store.put(("users", "1", "memories"), "food", {"text": "loves sushi", "kind": "pref"})
|
|
31
|
+
item = store.get(("users", "1", "memories"), "food")
|
|
32
|
+
hits = store.search(("users", "1"), filter={"kind": "pref"})
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
Use it as a LangGraph store: `graph.compile(store=FirestoreStore(...))`. Async methods work too (sync calls run in a thread). Authentication uses Application Default Credentials (`gcloud auth application-default login`).
|
|
36
|
+
|
|
37
|
+
## Semantic search (Firestore vector search)
|
|
38
|
+
|
|
39
|
+
Pass a LangGraph `IndexConfig` and the store embeds the configured fields on `put` (stored as a Firestore `Vector`) and ranks `search(query=...)` by cosine similarity using `find_nearest`:
|
|
40
|
+
|
|
41
|
+
```python
|
|
42
|
+
from langgraph_store_core import bedrock_titan_embeddings
|
|
43
|
+
from langgraph_store_firestore import FirestoreStore
|
|
44
|
+
|
|
45
|
+
store = FirestoreStore(project_id="my-gcp-project", collection="memory",
|
|
46
|
+
index={"dims": 1024, "embed": bedrock_titan_embeddings(dimensions=1024), "fields": ["text"]})
|
|
47
|
+
store.put(("memories", "kamal"), "k1", {"text": "the user loves sushi", "kind": "pref"})
|
|
48
|
+
hits = store.search(("memories", "kamal"), query="what food does the user like?", filter={"kind": "pref"})
|
|
49
|
+
print(hits[0].score, hits[0].value)
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
`embed` may be any LangChain `Embeddings`, a `list[str] -> list[list[float]]` callable, or a provider string. `fields` defaults to `["$"]` (whole value as JSON). `put(..., index=False)` skips embedding for one item.
|
|
53
|
+
|
|
54
|
+
Firestore requires a **vector index**. With `create_index=True` (default) the store creates a composite index `prefix ASC + embedding (flat, dims)` on the collection through the Firestore Admin API if one is missing; the build is asynchronous and `query` searches fail with `FAILED_PRECONDITION` until it is READY (usually a few minutes; check `gcloud firestore indexes composite list`). The caller needs `datastore.indexes.create` / `list` permissions, or create the index yourself and pass `create_index=False`:
|
|
55
|
+
|
|
56
|
+
```bash
|
|
57
|
+
gcloud firestore indexes composite create --collection-group=memory --query-scope=COLLECTION \
|
|
58
|
+
--field-config=order=ASCENDING,field-path=prefix \
|
|
59
|
+
--field-config=vector-config='{"dimension":"1024","flat":"{}"}',field-path=embedding
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
The namespace prefix is a range pre-filter on the vector query. Firestore pre-filters are limited to indexed equality/range fields, so value `filter`s are applied on the returned candidates (the store oversamples when a filter is given).
|
|
63
|
+
|
|
64
|
+
## Data model
|
|
65
|
+
|
|
66
|
+
A single collection of documents `{prefix, key, value, created_at, updated_at[, embedding: Vector]}`, doc id = `enc(prefix)__enc(key)`. Search is a `prefix` range query; filters and namespace matching are evaluated in the [core](https://pypi.org/project/langgraph-store-core/).
|
|
67
|
+
|
|
68
|
+
Docs: <https://skamalj.github.io/agentstate-reducer/>
|
|
69
|
+
|
|
70
|
+
## License
|
|
71
|
+
|
|
72
|
+
MIT
|
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
# langgraph-store-firestore
|
|
2
|
+
|
|
3
|
+
A **Google Firestore** long-term-memory store (`BaseStore`) for [LangGraph](https://langchain-ai.github.io/langgraph/) — namespaced key/value memory with prefix search, filters, `list_namespaces`, and **native semantic search via Firestore vector search**.
|
|
4
|
+
|
|
5
|
+
```bash
|
|
6
|
+
pip install langgraph-store-firestore
|
|
7
|
+
```
|
|
8
|
+
|
|
9
|
+
```python
|
|
10
|
+
from langgraph_store_firestore import FirestoreStore
|
|
11
|
+
|
|
12
|
+
store = FirestoreStore(project_id="my-gcp-project")
|
|
13
|
+
store.put(("users", "1", "memories"), "food", {"text": "loves sushi", "kind": "pref"})
|
|
14
|
+
item = store.get(("users", "1", "memories"), "food")
|
|
15
|
+
hits = store.search(("users", "1"), filter={"kind": "pref"})
|
|
16
|
+
```
|
|
17
|
+
|
|
18
|
+
Use it as a LangGraph store: `graph.compile(store=FirestoreStore(...))`. Async methods work too (sync calls run in a thread). Authentication uses Application Default Credentials (`gcloud auth application-default login`).
|
|
19
|
+
|
|
20
|
+
## Semantic search (Firestore vector search)
|
|
21
|
+
|
|
22
|
+
Pass a LangGraph `IndexConfig` and the store embeds the configured fields on `put` (stored as a Firestore `Vector`) and ranks `search(query=...)` by cosine similarity using `find_nearest`:
|
|
23
|
+
|
|
24
|
+
```python
|
|
25
|
+
from langgraph_store_core import bedrock_titan_embeddings
|
|
26
|
+
from langgraph_store_firestore import FirestoreStore
|
|
27
|
+
|
|
28
|
+
store = FirestoreStore(project_id="my-gcp-project", collection="memory",
|
|
29
|
+
index={"dims": 1024, "embed": bedrock_titan_embeddings(dimensions=1024), "fields": ["text"]})
|
|
30
|
+
store.put(("memories", "kamal"), "k1", {"text": "the user loves sushi", "kind": "pref"})
|
|
31
|
+
hits = store.search(("memories", "kamal"), query="what food does the user like?", filter={"kind": "pref"})
|
|
32
|
+
print(hits[0].score, hits[0].value)
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
`embed` may be any LangChain `Embeddings`, a `list[str] -> list[list[float]]` callable, or a provider string. `fields` defaults to `["$"]` (whole value as JSON). `put(..., index=False)` skips embedding for one item.
|
|
36
|
+
|
|
37
|
+
Firestore requires a **vector index**. With `create_index=True` (default) the store creates a composite index `prefix ASC + embedding (flat, dims)` on the collection through the Firestore Admin API if one is missing; the build is asynchronous and `query` searches fail with `FAILED_PRECONDITION` until it is READY (usually a few minutes; check `gcloud firestore indexes composite list`). The caller needs `datastore.indexes.create` / `list` permissions, or create the index yourself and pass `create_index=False`:
|
|
38
|
+
|
|
39
|
+
```bash
|
|
40
|
+
gcloud firestore indexes composite create --collection-group=memory --query-scope=COLLECTION \
|
|
41
|
+
--field-config=order=ASCENDING,field-path=prefix \
|
|
42
|
+
--field-config=vector-config='{"dimension":"1024","flat":"{}"}',field-path=embedding
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
The namespace prefix is a range pre-filter on the vector query. Firestore pre-filters are limited to indexed equality/range fields, so value `filter`s are applied on the returned candidates (the store oversamples when a filter is given).
|
|
46
|
+
|
|
47
|
+
## Data model
|
|
48
|
+
|
|
49
|
+
A single collection of documents `{prefix, key, value, created_at, updated_at[, embedding: Vector]}`, doc id = `enc(prefix)__enc(key)`. Search is a `prefix` range query; filters and namespace matching are evaluated in the [core](https://pypi.org/project/langgraph-store-core/).
|
|
50
|
+
|
|
51
|
+
Docs: <https://skamalj.github.io/agentstate-reducer/>
|
|
52
|
+
|
|
53
|
+
## License
|
|
54
|
+
|
|
55
|
+
MIT
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "langgraph-store-firestore"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Google Firestore long-term-memory store (BaseStore) for LangGraph — namespaced key/value memory with filters, list_namespaces and native semantic (Firestore vector) search"
|
|
9
|
+
authors = [{name = "Kamal", email = "skamalj@gmail.com"}]
|
|
10
|
+
readme = "README.md"
|
|
11
|
+
requires-python = ">=3.10"
|
|
12
|
+
dependencies = [
|
|
13
|
+
"langgraph-store-core>=0.1.0",
|
|
14
|
+
"google-cloud-firestore",
|
|
15
|
+
]
|
|
16
|
+
keywords = ["langgraph", "store", "basestore", "firestore", "gcp", "long-term-memory", "agent-memory"]
|
|
17
|
+
classifiers = [
|
|
18
|
+
"Programming Language :: Python :: 3",
|
|
19
|
+
"License :: OSI Approved :: MIT License",
|
|
20
|
+
"Operating System :: OS Independent",
|
|
21
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
[project.urls]
|
|
25
|
+
Homepage = "https://github.com/skamalj/langgraph-store"
|
|
26
|
+
Repository = "https://github.com/skamalj/langgraph-store.git"
|
|
27
|
+
|
|
28
|
+
[dependency-groups]
|
|
29
|
+
dev = ["pytest>=7.0", "pytest-asyncio"]
|
|
30
|
+
|
|
31
|
+
[tool.pytest.ini_options]
|
|
32
|
+
asyncio_mode = "auto"
|
|
33
|
+
|
|
34
|
+
[tool.hatch.build.targets.wheel]
|
|
35
|
+
packages = ["src/langgraph_store_firestore"]
|
|
@@ -0,0 +1,171 @@
|
|
|
1
|
+
"""Google Firestore implementation of the LangGraph long-term-memory ``BaseStore``.
|
|
2
|
+
|
|
3
|
+
``FirestoreStore`` persists namespaced key/value memory in a single collection and
|
|
4
|
+
supports get / put / delete / search (namespace-prefix + filter) and
|
|
5
|
+
``list_namespaces``. With an ``IndexConfig`` it also does **semantic search**
|
|
6
|
+
natively via Firestore vector search (``find_nearest``, cosine). Built on
|
|
7
|
+
``langgraph-store-core``.
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
import logging
|
|
13
|
+
from typing import List, Optional, Tuple
|
|
14
|
+
from urllib.parse import quote
|
|
15
|
+
|
|
16
|
+
from google.cloud import firestore
|
|
17
|
+
from google.cloud.firestore_v1.base_query import FieldFilter
|
|
18
|
+
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
|
|
19
|
+
from google.cloud.firestore_v1.vector import Vector
|
|
20
|
+
|
|
21
|
+
from langgraph_store_core import IndexConfig, KVStore
|
|
22
|
+
|
|
23
|
+
__all__ = ["FirestoreStore"]
|
|
24
|
+
|
|
25
|
+
logger = logging.getLogger(__name__)
|
|
26
|
+
|
|
27
|
+
_FIELDS = ("prefix", "key", "value", "created_at", "updated_at")
|
|
28
|
+
_PREFIX_HIGH = chr(0xF8FF) # exclusive upper bound for a Firestore prefix scan
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
class FirestoreStore(KVStore):
|
|
32
|
+
"""LangGraph ``BaseStore`` backed by Google Firestore.
|
|
33
|
+
|
|
34
|
+
Example:
|
|
35
|
+
```python
|
|
36
|
+
from langgraph_store_firestore import FirestoreStore
|
|
37
|
+
|
|
38
|
+
store = FirestoreStore(project_id="my-gcp-project")
|
|
39
|
+
store.put(("users", "1"), "profile", {"name": "Kamal"})
|
|
40
|
+
|
|
41
|
+
# semantic search
|
|
42
|
+
store = FirestoreStore(project_id=..., index={"dims": 1024, "embed": my_embedder, "fields": ["text"]})
|
|
43
|
+
hits = store.search(("users", "1"), query="what does the user like?")
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
Authentication uses Application Default Credentials. With ``index`` the
|
|
47
|
+
embedding is stored as a Firestore ``Vector`` on ``embedding`` and queries
|
|
48
|
+
use ``find_nearest`` (cosine). Firestore requires a **vector index** on the
|
|
49
|
+
collection; when ``create_index=True`` (default) the store creates a
|
|
50
|
+
composite index ``prefix ASC + embedding (flat, dims)`` via the Admin API if
|
|
51
|
+
one is missing (index build is asynchronous; the first queries may fail
|
|
52
|
+
until it is READY).
|
|
53
|
+
"""
|
|
54
|
+
|
|
55
|
+
supports_native_vector_search = True
|
|
56
|
+
|
|
57
|
+
def __init__(
|
|
58
|
+
self,
|
|
59
|
+
*,
|
|
60
|
+
project_id: str,
|
|
61
|
+
collection: str = "langgraph_store",
|
|
62
|
+
client: Optional["firestore.Client"] = None,
|
|
63
|
+
index: Optional[IndexConfig] = None,
|
|
64
|
+
create_index: bool = True,
|
|
65
|
+
database: str = "(default)",
|
|
66
|
+
) -> None:
|
|
67
|
+
super().__init__(index=index)
|
|
68
|
+
self._project = project_id
|
|
69
|
+
self._database = database
|
|
70
|
+
self._client = client or firestore.Client(project=project_id, database=database)
|
|
71
|
+
self._col = self._client.collection(collection)
|
|
72
|
+
self._collection_name = collection
|
|
73
|
+
if self.dims and create_index:
|
|
74
|
+
self._ensure_vector_index()
|
|
75
|
+
|
|
76
|
+
# ------------------------------------------------------------------ index bootstrap
|
|
77
|
+
def _ensure_vector_index(self) -> None:
|
|
78
|
+
"""Create the composite vector index (prefix ASC, embedding flat) if absent."""
|
|
79
|
+
from google.cloud import firestore_admin_v1 as fa
|
|
80
|
+
|
|
81
|
+
admin = fa.FirestoreAdminClient()
|
|
82
|
+
parent = (
|
|
83
|
+
f"projects/{self._project}/databases/{self._database}"
|
|
84
|
+
f"/collectionGroups/{self._collection_name}"
|
|
85
|
+
)
|
|
86
|
+
for ix in admin.list_indexes(parent=parent):
|
|
87
|
+
paths = [f.field_path for f in ix.fields]
|
|
88
|
+
if "embedding" in paths and "prefix" in paths:
|
|
89
|
+
return
|
|
90
|
+
index = fa.Index(
|
|
91
|
+
query_scope=fa.Index.QueryScope.COLLECTION,
|
|
92
|
+
fields=[
|
|
93
|
+
fa.Index.IndexField(field_path="prefix", order=fa.Index.IndexField.Order.ASCENDING),
|
|
94
|
+
fa.Index.IndexField(
|
|
95
|
+
field_path="embedding",
|
|
96
|
+
vector_config=fa.Index.IndexField.VectorConfig(
|
|
97
|
+
dimension=self.dims, flat=fa.Index.IndexField.VectorConfig.FlatIndex()
|
|
98
|
+
),
|
|
99
|
+
),
|
|
100
|
+
],
|
|
101
|
+
)
|
|
102
|
+
try:
|
|
103
|
+
admin.create_index(parent=parent, index=index)
|
|
104
|
+
logger.info("FirestoreStore: creating vector index on %s (async)", self._collection_name)
|
|
105
|
+
except Exception as exc: # already exists / racing creation
|
|
106
|
+
logger.info("FirestoreStore: vector index create skipped: %s", exc)
|
|
107
|
+
|
|
108
|
+
# ------------------------------------------------------------------ primitives
|
|
109
|
+
@staticmethod
|
|
110
|
+
def _docid(prefix: str, key: str) -> str:
|
|
111
|
+
# Firestore doc ids can't contain '/'; url-encode both parts, join safely.
|
|
112
|
+
return f"{quote(prefix, safe='')}__{quote(key, safe='')}"
|
|
113
|
+
|
|
114
|
+
def _row(self, d: dict) -> dict:
|
|
115
|
+
row = {f: d[f] for f in _FIELDS}
|
|
116
|
+
emb = d.get("embedding")
|
|
117
|
+
row["embedding"] = list(emb.to_map_value()["value"]) if isinstance(emb, Vector) else (
|
|
118
|
+
list(emb) if emb is not None else None
|
|
119
|
+
)
|
|
120
|
+
return row
|
|
121
|
+
|
|
122
|
+
def _read(self, prefix: str, key: str) -> Optional[dict]:
|
|
123
|
+
doc = self._col.document(self._docid(prefix, key)).get()
|
|
124
|
+
if not doc.exists:
|
|
125
|
+
return None
|
|
126
|
+
return self._row(doc.to_dict())
|
|
127
|
+
|
|
128
|
+
def _write(self, row: dict) -> None:
|
|
129
|
+
doc = {f: row[f] for f in _FIELDS}
|
|
130
|
+
if self.dims and row.get("embedding") is not None:
|
|
131
|
+
doc["embedding"] = Vector(row["embedding"])
|
|
132
|
+
self._col.document(self._docid(row["prefix"], row["key"])).set(doc)
|
|
133
|
+
|
|
134
|
+
def _remove(self, prefix: str, key: str) -> None:
|
|
135
|
+
self._col.document(self._docid(prefix, key)).delete()
|
|
136
|
+
|
|
137
|
+
def _prefix_query(self, prefix: str):
|
|
138
|
+
if prefix == "":
|
|
139
|
+
return self._col
|
|
140
|
+
return (
|
|
141
|
+
self._col.where(filter=FieldFilter("prefix", ">=", prefix))
|
|
142
|
+
.where(filter=FieldFilter("prefix", "<", prefix + _PREFIX_HIGH))
|
|
143
|
+
)
|
|
144
|
+
|
|
145
|
+
def _scan(self, prefix: str) -> List[dict]:
|
|
146
|
+
return [self._row(d.to_dict()) for d in self._prefix_query(prefix).stream()]
|
|
147
|
+
|
|
148
|
+
def _vector_search(
|
|
149
|
+
self, prefix: str, vector: List[float], filter: Optional[dict], limit: int
|
|
150
|
+
) -> List[Tuple[dict, float]]:
|
|
151
|
+
"""Native Firestore ``find_nearest`` (cosine) after a prefix range pre-filter.
|
|
152
|
+
|
|
153
|
+
Firestore pre-filters are limited to indexed equality/range fields, so
|
|
154
|
+
value filters are applied by core on the returned candidates; we
|
|
155
|
+
oversample to leave room for that.
|
|
156
|
+
"""
|
|
157
|
+
oversample = limit * 4 if filter else limit
|
|
158
|
+
q = self._prefix_query(prefix).find_nearest(
|
|
159
|
+
vector_field="embedding",
|
|
160
|
+
query_vector=Vector(vector),
|
|
161
|
+
distance_measure=DistanceMeasure.COSINE,
|
|
162
|
+
limit=oversample,
|
|
163
|
+
distance_result_field="_distance",
|
|
164
|
+
)
|
|
165
|
+
out: List[Tuple[dict, float]] = []
|
|
166
|
+
for d in q.stream():
|
|
167
|
+
data = d.to_dict()
|
|
168
|
+
dist = data.pop("_distance", None)
|
|
169
|
+
row = self._row(data)
|
|
170
|
+
out.append((row, 1.0 - float(dist) if dist is not None else 0.0))
|
|
171
|
+
return out
|
|
@@ -0,0 +1,118 @@
|
|
|
1
|
+
"""Semantic-search tests for FirestoreStore (native find_nearest path) against a real Firestore database.
|
|
2
|
+
|
|
3
|
+
The store creates the composite vector index (prefix ASC + embedding flat) via the
|
|
4
|
+
Admin API. Index builds are asynchronous, so the fixture waits for READY, and the
|
|
5
|
+
collection name is fixed per DIMS so the index is reused across runs.
|
|
6
|
+
"""
|
|
7
|
+
import os
|
|
8
|
+
import time
|
|
9
|
+
import uuid
|
|
10
|
+
|
|
11
|
+
import pytest
|
|
12
|
+
|
|
13
|
+
from langgraph_store_core.testing import FakeEmbeddings
|
|
14
|
+
from langgraph_store_firestore import FirestoreStore
|
|
15
|
+
|
|
16
|
+
PROJECT = (
|
|
17
|
+
os.environ.get("PAI_FIRESTORE_PROJECT")
|
|
18
|
+
or os.environ.get("GCP_PROJECT")
|
|
19
|
+
or os.environ.get("GOOGLE_CLOUD_PROJECT")
|
|
20
|
+
or "gcdeveloper-new"
|
|
21
|
+
)
|
|
22
|
+
DIMS = 64
|
|
23
|
+
COLLECTION = f"lg_vstore_test_d{DIMS}" # stable name so the vector index persists between runs
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def _fs_available() -> bool:
|
|
27
|
+
try:
|
|
28
|
+
from google.cloud import firestore
|
|
29
|
+
firestore.Client(project=PROJECT).collection("x").limit(1).get()
|
|
30
|
+
return True
|
|
31
|
+
except Exception:
|
|
32
|
+
return False
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
pytestmark = pytest.mark.skipif(not _fs_available(), reason="Firestore not reachable")
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _wait_index_ready(project: str, collection: str, timeout: int = 600) -> None:
|
|
39
|
+
from google.cloud import firestore_admin_v1 as fa
|
|
40
|
+
admin = fa.FirestoreAdminClient()
|
|
41
|
+
parent = f"projects/{project}/databases/(default)/collectionGroups/{collection}"
|
|
42
|
+
deadline = time.time() + timeout
|
|
43
|
+
while time.time() < deadline:
|
|
44
|
+
for ix in admin.list_indexes(parent=parent):
|
|
45
|
+
paths = [f.field_path for f in ix.fields]
|
|
46
|
+
if "embedding" in paths and ix.state == fa.Index.State.READY:
|
|
47
|
+
return
|
|
48
|
+
time.sleep(10)
|
|
49
|
+
raise RuntimeError("vector index not READY in time")
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
@pytest.fixture(scope="module")
|
|
53
|
+
def store():
|
|
54
|
+
s = FirestoreStore(project_id=PROJECT, collection=COLLECTION,
|
|
55
|
+
index={"dims": DIMS, "embed": FakeEmbeddings(DIMS), "fields": ["text"]})
|
|
56
|
+
_wait_index_ready(PROJECT, COLLECTION)
|
|
57
|
+
yield s
|
|
58
|
+
for row in s._scan(""):
|
|
59
|
+
s._remove(row["prefix"], row["key"])
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
@pytest.fixture(autouse=True)
|
|
63
|
+
def _clean(store):
|
|
64
|
+
yield
|
|
65
|
+
for row in store._scan(""):
|
|
66
|
+
store._remove(row["prefix"], row["key"])
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def _seed(s):
|
|
70
|
+
s.put(("memories", "kamal"), "food", {"text": "the user loves sushi and japanese food", "kind": "pref"})
|
|
71
|
+
s.put(("memories", "kamal"), "lang", {"text": "the user writes python every day", "kind": "fact"})
|
|
72
|
+
s.put(("memories", "kamal"), "city", {"text": "the user lives by the sea in a big city", "kind": "fact"})
|
|
73
|
+
s.put(("memories", "priya"), "food", {"text": "loves sushi", "kind": "pref"})
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def test_put_embeds_and_index_false_skips(store):
|
|
77
|
+
store.put(("n",), "k", {"text": "hello world"})
|
|
78
|
+
assert len(store._read("n", "k")["embedding"]) == DIMS
|
|
79
|
+
store.put(("n",), "k2", {"text": "hello"}, index=False)
|
|
80
|
+
assert store._read("n", "k2")["embedding"] is None
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def test_query_ranks_and_scores(store):
|
|
84
|
+
_seed(store)
|
|
85
|
+
hits = store.search(("memories", "kamal"), query="sushi and japanese food")
|
|
86
|
+
assert hits[0].key == "food" and hits[0].score is not None
|
|
87
|
+
assert hits[0].score >= hits[-1].score
|
|
88
|
+
assert all(h.namespace == ("memories", "kamal") for h in hits)
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
def test_query_prefix_children_and_filter(store):
|
|
92
|
+
_seed(store)
|
|
93
|
+
keys = {(h.namespace, h.key) for h in store.search(("memories",), query="sushi")}
|
|
94
|
+
assert (("memories", "kamal"), "food") in keys and (("memories", "priya"), "food") in keys
|
|
95
|
+
facts = store.search(("memories", "kamal"), query="user", filter={"kind": "fact"})
|
|
96
|
+
assert facts and all(h.value["kind"] == "fact" for h in facts)
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
def test_query_limit_offset(store):
|
|
100
|
+
_seed(store)
|
|
101
|
+
first = store.search(("memories", "kamal"), query="user", limit=2)
|
|
102
|
+
rest = store.search(("memories", "kamal"), query="user", limit=2, offset=2)
|
|
103
|
+
assert len(first) == 2 and len(rest) == 1 and not {h.key for h in first} & {h.key for h in rest}
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def test_no_query_is_filter_only(store):
|
|
107
|
+
_seed(store)
|
|
108
|
+
hits = store.search(("memories", "kamal"))
|
|
109
|
+
assert [h.key for h in hits] == ["city", "food", "lang"] and all(h.score is None for h in hits)
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
def test_native_path_used(store, monkeypatch):
|
|
113
|
+
_seed(store)
|
|
114
|
+
called = []
|
|
115
|
+
orig = store._vector_search
|
|
116
|
+
monkeypatch.setattr(store, "_vector_search", lambda *a, **k: (called.append(1), orig(*a, **k))[1])
|
|
117
|
+
store.search(("memories",), query="sushi")
|
|
118
|
+
assert called
|
|
@@ -0,0 +1,85 @@
|
|
|
1
|
+
"""Tests for FirestoreStore (LangGraph BaseStore) against real Firestore.
|
|
2
|
+
|
|
3
|
+
Requires a GCP project (PAI_FIRESTORE_PROJECT / GCP_PROJECT / GOOGLE_CLOUD_PROJECT)
|
|
4
|
+
and Application Default Credentials. A unique collection per test isolates state.
|
|
5
|
+
"""
|
|
6
|
+
import os
|
|
7
|
+
import uuid
|
|
8
|
+
|
|
9
|
+
import pytest
|
|
10
|
+
|
|
11
|
+
from langgraph_store_firestore import FirestoreStore
|
|
12
|
+
|
|
13
|
+
PROJECT = (
|
|
14
|
+
os.environ.get("PAI_FIRESTORE_PROJECT")
|
|
15
|
+
or os.environ.get("GCP_PROJECT")
|
|
16
|
+
or os.environ.get("GOOGLE_CLOUD_PROJECT")
|
|
17
|
+
)
|
|
18
|
+
|
|
19
|
+
pytestmark = pytest.mark.skipif(not PROJECT, reason="no GCP project env var set")
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
@pytest.fixture()
|
|
23
|
+
def store():
|
|
24
|
+
cname = f"lg_store_test_{uuid.uuid4().hex[:8]}"
|
|
25
|
+
s = FirestoreStore(project_id=PROJECT, collection=cname)
|
|
26
|
+
yield s
|
|
27
|
+
for d in s._col.stream():
|
|
28
|
+
d.reference.delete()
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def test_put_get(store):
|
|
32
|
+
store.put(("users", "1"), "food", {"text": "sushi"})
|
|
33
|
+
it = store.get(("users", "1"), "food")
|
|
34
|
+
assert it is not None and it.value == {"text": "sushi"}
|
|
35
|
+
assert it.namespace == ("users", "1") and it.key == "food"
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def test_get_missing_is_none(store):
|
|
39
|
+
assert store.get(("nope",), "x") is None
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def test_update_preserves_created_at(store):
|
|
43
|
+
store.put(("u",), "k", {"v": 1})
|
|
44
|
+
a = store.get(("u",), "k")
|
|
45
|
+
store.put(("u",), "k", {"v": 2})
|
|
46
|
+
b = store.get(("u",), "k")
|
|
47
|
+
assert b.value == {"v": 2}
|
|
48
|
+
assert b.created_at == a.created_at
|
|
49
|
+
assert b.updated_at >= a.updated_at
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def test_delete_both_ways(store):
|
|
53
|
+
store.put(("u",), "k", {"v": 1})
|
|
54
|
+
store.delete(("u",), "k")
|
|
55
|
+
assert store.get(("u",), "k") is None
|
|
56
|
+
store.put(("u",), "k2", {"v": 1})
|
|
57
|
+
store.put(("u",), "k2", None)
|
|
58
|
+
assert store.get(("u",), "k2") is None
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def test_search_prefix_and_filter(store):
|
|
62
|
+
store.put(("users", "1", "mem"), "a", {"kind": "pref"})
|
|
63
|
+
store.put(("users", "1", "mem"), "b", {"kind": "fact"})
|
|
64
|
+
store.put(("users", "2", "mem"), "c", {"kind": "pref"})
|
|
65
|
+
assert sorted(i.key for i in store.search(("users", "1"))) == ["a", "b"]
|
|
66
|
+
assert [i.key for i in store.search(("users", "1"), filter={"kind": "pref"})] == ["a"]
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def test_list_namespaces(store):
|
|
70
|
+
store.put(("users", "1", "mem"), "a", {})
|
|
71
|
+
store.put(("users", "2", "mem"), "b", {})
|
|
72
|
+
store.put(("orgs", "x"), "c", {})
|
|
73
|
+
allns = store.list_namespaces()
|
|
74
|
+
assert ("users", "1", "mem") in allns and ("orgs", "x") in allns
|
|
75
|
+
assert all(n[0] == "users" for n in store.list_namespaces(prefix=("users",)))
|
|
76
|
+
depth1 = store.list_namespaces(max_depth=1)
|
|
77
|
+
assert ("users",) in depth1 and ("orgs",) in depth1
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
async def test_async_surface(store):
|
|
81
|
+
await store.aput(("u",), "k", {"v": 1})
|
|
82
|
+
it = await store.aget(("u",), "k")
|
|
83
|
+
assert it.value == {"v": 1}
|
|
84
|
+
await store.adelete(("u",), "k")
|
|
85
|
+
assert await store.aget(("u",), "k") is None
|