fix: add auto_refresh option for OpenSearch Serverless compatibility (#3893)
Co-authored-by: Zaid Malhis <Malhis@users.noreply.github.com>
This commit is contained in:
@@ -56,6 +56,30 @@ config = {
|
|||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
|
### Configuration Options
|
||||||
|
|
||||||
|
| Parameter | Type | Default | Description |
|
||||||
|
|-----------|------|---------|-------------|
|
||||||
|
| `collection_name` | string | required | Name of the OpenSearch index |
|
||||||
|
| `host` | string | required | OpenSearch endpoint URL |
|
||||||
|
| `port` | int | 9200 | Port number |
|
||||||
|
| `http_auth` | object | None | Authentication credentials (e.g., AWSV4SignerAuth) |
|
||||||
|
| `embedding_model_dims` | int | 1536 | Dimension of embedding vectors |
|
||||||
|
| `use_ssl` | bool | False | Enable SSL/TLS connection |
|
||||||
|
| `verify_certs` | bool | False | Verify SSL certificates |
|
||||||
|
| `auto_refresh` | bool | False | Automatically refresh index after insert. OpenSearch refreshes every ~1 second by default, so this is rarely needed. |
|
||||||
|
|
||||||
|
<Note>
|
||||||
|
The defaults above match a local OpenSearch instance. The AWS OpenSearch Serverless
|
||||||
|
example earlier on this page intentionally overrides them with `port=443`, `use_ssl=True`,
|
||||||
|
and `verify_certs=True`, which are required when connecting to a Serverless collection.
|
||||||
|
</Note>
|
||||||
|
|
||||||
|
<Note>
|
||||||
|
For **AWS OpenSearch Serverless**, keep `auto_refresh=False` (the default).
|
||||||
|
The `indices.refresh()` API is not supported on Serverless collections.
|
||||||
|
</Note>
|
||||||
|
|
||||||
### Add Memories
|
### Add Memories
|
||||||
|
|
||||||
```python
|
```python
|
||||||
|
|||||||
@@ -18,6 +18,12 @@ class OpenSearchConfig(BaseModel):
|
|||||||
"RequestsHttpConnection", description="Connection class for OpenSearch"
|
"RequestsHttpConnection", description="Connection class for OpenSearch"
|
||||||
)
|
)
|
||||||
pool_maxsize: int = Field(20, description="Maximum number of connections in the pool")
|
pool_maxsize: int = Field(20, description="Maximum number of connections in the pool")
|
||||||
|
auto_refresh: bool = Field(
|
||||||
|
False,
|
||||||
|
description="Automatically refresh index after insert operations to make documents "
|
||||||
|
"immediately searchable. Disabled by default for OpenSearch Serverless compatibility. "
|
||||||
|
"OpenSearch automatically refreshes indices every ~1 second, so most users don't need this.",
|
||||||
|
)
|
||||||
|
|
||||||
@model_validator(mode="before")
|
@model_validator(mode="before")
|
||||||
@classmethod
|
@classmethod
|
||||||
|
|||||||
@@ -39,6 +39,8 @@ class OpenSearchDB(VectorStoreBase):
|
|||||||
|
|
||||||
self.collection_name = config.collection_name
|
self.collection_name = config.collection_name
|
||||||
self.embedding_model_dims = config.embedding_model_dims
|
self.embedding_model_dims = config.embedding_model_dims
|
||||||
|
self.auto_refresh = config.auto_refresh
|
||||||
|
|
||||||
self.create_col(self.collection_name, self.embedding_model_dims)
|
self.create_col(self.collection_name, self.embedding_model_dims)
|
||||||
|
|
||||||
def create_index(self) -> None:
|
def create_index(self) -> None:
|
||||||
@@ -148,8 +150,6 @@ class OpenSearchDB(VectorStoreBase):
|
|||||||
}
|
}
|
||||||
try:
|
try:
|
||||||
self.client.index(index=self.collection_name, body=body)
|
self.client.index(index=self.collection_name, body=body)
|
||||||
# Force refresh to make documents immediately searchable for tests
|
|
||||||
self.client.indices.refresh(index=self.collection_name)
|
|
||||||
|
|
||||||
results.append(
|
results.append(
|
||||||
OutputData(
|
OutputData(
|
||||||
@@ -162,6 +162,14 @@ class OpenSearchDB(VectorStoreBase):
|
|||||||
logger.error(f"Error inserting vector {id_}: {e}", exc_info=True)
|
logger.error(f"Error inserting vector {id_}: {e}", exc_info=True)
|
||||||
raise
|
raise
|
||||||
|
|
||||||
|
# Refresh once after the full batch (not per document) if explicitly enabled.
|
||||||
|
# Disabled by default for Serverless compatibility: OpenSearch Serverless does not
|
||||||
|
# support the indices.refresh() API, and refreshing per document would cause a
|
||||||
|
# cluster-level I/O stall on every insert.
|
||||||
|
# See: https://docs.aws.amazon.com/opensearch-service/latest/developerguide/serverless-genref.html
|
||||||
|
if self.auto_refresh:
|
||||||
|
self.client.indices.refresh(index=self.collection_name)
|
||||||
|
|
||||||
return results
|
return results
|
||||||
|
|
||||||
def search(
|
def search(
|
||||||
|
|||||||
@@ -151,6 +151,55 @@ class TestOpenSearchDB(unittest.TestCase):
|
|||||||
self.os_db.create_index()
|
self.os_db.create_index()
|
||||||
self.client_mock.indices.create.assert_not_called()
|
self.client_mock.indices.create.assert_not_called()
|
||||||
|
|
||||||
|
def test_auto_refresh_disabled_by_default(self):
|
||||||
|
"""Test that auto_refresh is disabled by default (Issue #3739).
|
||||||
|
|
||||||
|
This ensures OpenSearch Serverless compatibility out-of-the-box since
|
||||||
|
the indices.refresh() API is not supported in serverless mode.
|
||||||
|
"""
|
||||||
|
# Default instance should have auto_refresh=False
|
||||||
|
self.assertFalse(self.os_db.auto_refresh)
|
||||||
|
self.client_mock.reset_mock()
|
||||||
|
|
||||||
|
vectors = [[0.1] * 1536]
|
||||||
|
payloads = [{"key1": "value1"}]
|
||||||
|
ids = ["id1"]
|
||||||
|
|
||||||
|
self.os_db.insert(vectors=vectors, payloads=payloads, ids=ids)
|
||||||
|
|
||||||
|
# Verify index was called but refresh was NOT called (default behavior)
|
||||||
|
self.assertEqual(self.client_mock.index.call_count, 1)
|
||||||
|
self.client_mock.indices.refresh.assert_not_called()
|
||||||
|
|
||||||
|
def test_auto_refresh_enabled(self):
|
||||||
|
"""Test that refresh is called once per batch (not per document) when auto_refresh=True."""
|
||||||
|
with patch("mem0.vector_stores.opensearch.OpenSearch", return_value=self.client_mock):
|
||||||
|
auto_refresh_db = OpenSearchDB(
|
||||||
|
host="localhost",
|
||||||
|
port=9200,
|
||||||
|
collection_name="test_auto_refresh",
|
||||||
|
embedding_model_dims=1536,
|
||||||
|
auto_refresh=True, # Enable auto-refresh
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertTrue(auto_refresh_db.auto_refresh)
|
||||||
|
# auto_refresh_db reuses self.client_mock (patched above), so reset to drop
|
||||||
|
# the index calls made during construction before asserting on insert().
|
||||||
|
self.client_mock.reset_mock()
|
||||||
|
|
||||||
|
# Insert a batch of 3 vectors to verify the refresh is hoisted out of the
|
||||||
|
# per-document loop: index() is called once per document, but refresh()
|
||||||
|
# must fire exactly once for the whole batch.
|
||||||
|
vectors = [[0.1] * 1536, [0.2] * 1536, [0.3] * 1536]
|
||||||
|
payloads = [{"key1": "value1"}, {"key2": "value2"}, {"key3": "value3"}]
|
||||||
|
ids = ["id1", "id2", "id3"]
|
||||||
|
|
||||||
|
auto_refresh_db.insert(vectors=vectors, payloads=payloads, ids=ids)
|
||||||
|
|
||||||
|
# index() once per document, but refresh() only once for the batch
|
||||||
|
self.assertEqual(self.client_mock.index.call_count, 3)
|
||||||
|
self.assertEqual(self.client_mock.indices.refresh.call_count, 1)
|
||||||
|
|
||||||
def test_insert(self):
|
def test_insert(self):
|
||||||
vectors = [[0.1] * 1536, [0.2] * 1536]
|
vectors = [[0.1] * 1536, [0.2] * 1536]
|
||||||
payloads = [{"key1": "value1"}, {"key2": "value2"}]
|
payloads = [{"key1": "value1"}, {"key2": "value2"}]
|
||||||
|
|||||||
Reference in New Issue
Block a user