wehub/hkuds--lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
a172f1251d4fa807bf948d97bb0f9a2e2bcb63d2
hkuds--lightrag/lightrag/kg/opensearch_impl.py
T
wehub-resource-sync 4a4a1fed67 chore: import upstream snapshot with attribution
2026-07-13 12:08:54 +08:00

4588 lines
191 KiB
Python
Raw Blame History

"""
OpenSearch Storage Implementation for LightRAG
This module provides OpenSearch-based storage backends for LightRAG,
including KV storage, document status storage, graph storage, and vector storage.
Requirements:
- opensearch-py >= 3.0.0
- OpenSearch 3.x or higher with k-NN plugin enabled
"""
import os
import re
import ssl as ssl_module
import time
import asyncio
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Union, final
import numpy as np
import configparser
from ..base import (
BaseGraphStorage,
BaseKVStorage,
BaseVectorStorage,
DocProcessingStatus,
DocStatus,
DocStatusStorage,
)
from ..utils import (
logger,
compute_mdhash_id,
_cooperative_yield,
merge_source_ids,
validate_workspace,
)
from ..types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge
from ..constants import GRAPH_FIELD_SEP, DEFAULT_QUERY_PRIORITY
from ..kg.shared_storage import get_data_init_lock, get_namespace_lock
import pipmaster as pm
if not pm.is_installed("opensearch-py"):
pm.install("opensearch-py")
from opensearchpy import AsyncOpenSearch, helpers # type: ignore
from opensearchpy.exceptions import ( # type: ignore
OpenSearchException,
NotFoundError,
RequestError,
ConflictError,
)
config = configparser.ConfigParser()
config.read("config.ini", "utf-8")
def _get_opensearch_env(key, fallback):
cfg_key = key.replace("OPENSEARCH_", "").lower()
return os.environ.get(key, config.get("opensearch", cfg_key, fallback=fallback))
def _get_index_number_of_shards() -> int:
return int(_get_opensearch_env("OPENSEARCH_NUMBER_OF_SHARDS", "1"))
def _get_index_number_of_replicas() -> int:
return int(_get_opensearch_env("OPENSEARCH_NUMBER_OF_REPLICAS", "0"))
def _sanitize_index_name(name: str) -> str:
"""Sanitize a string to be a valid OpenSearch index name."""
sanitized = re.sub(r"[^a-z0-9_-]", "_", name.lower())
if sanitized and sanitized[0] in "-_+":
sanitized = "x" + sanitized
return sanitized
# HTTP statuses that indicate a transient failure where retrying makes sense:
# request timeout, rate limit, and the standard 5xx server-error range.
# A missing status (None) typically means a network or parse error before the
# server responded, which is also retriable.
_RETRYABLE_BULK_STATUSES: frozenset[int] = frozenset({408, 429, 500, 502, 503, 504})
# Cap the length of error summaries dumped to logs so a multi-MB mapping
# explanation can't flood the log file.
_BULK_ERROR_SUMMARY_MAX_LEN = 200
@dataclass(frozen=True)
class _FailedBulkOp:
"""Structured representation of a non-retryable per-action bulk failure."""
op: str
doc_id: str
status: int | None
error: str
@dataclass
class _PendingVectorDoc:
"""Buffered vector upsert waiting for embedding and/or bulk flush."""
source: dict[str, Any]
content: str
vector: list[float] | None = None
def _summarize_bulk_error(error: Any) -> str:
"""Turn an opensearch-py per-action ``error`` payload into a short string.
The field may be a string, dict (``{"type": ..., "reason": ...}``) or
something else entirely. We prefer ``reason`` / ``type`` from dicts to
keep the log readable.
"""
if error is None:
return ""
if isinstance(error, str):
summary = error
elif isinstance(error, dict):
reason = error.get("reason") or error.get("type")
summary = reason if isinstance(reason, str) else repr(error)
else:
summary = repr(error)
if len(summary) > _BULK_ERROR_SUMMARY_MAX_LEN:
summary = summary[: _BULK_ERROR_SUMMARY_MAX_LEN - 3] + "..."
return summary
def _extract_bulk_failed_ids(
failed: list[Any] | None,
) -> tuple[set[str], list[_FailedBulkOp]]:
"""Split an opensearch-py bulk ``failed`` list into retryable / dead ops.
``async_bulk(raise_on_error=False)`` returns ``(success, failed)`` where
``failed`` is a list of per-action error dicts shaped like::
{"index": {"_id": "...", "status": 500, "error": {...}}}
{"delete": {"_id": "...", "status": 404, ...}}
{"create": {"_id": "...", "status": 409, ...}}
Returns ``(retryable, non_retryable)``:
* ``retryable`` — ``set[str]`` of ids that should be retried on
the next flush (408 / 429 / 5xx, plus a missing status which
usually means a network-level failure before the server responded).
* ``non_retryable`` — ``list[_FailedBulkOp]`` of permanent failures
(most 4xx, mapping errors, etc.) carrying op-name, id, status and
a short ``error`` summary so callers can log meaningful context.
``404`` on a delete is treated as success-equivalent and dropped
from both sets.
Unrecognised or malformed entries are skipped so a stray dict shape
never crashes the flush path.
"""
retryable: set[str] = set()
non_retryable: list[_FailedBulkOp] = []
if not failed:
return retryable, non_retryable
for entry in failed:
if not isinstance(entry, dict):
continue
for op_name, op_payload in entry.items():
if not isinstance(op_payload, dict):
continue
doc_id = op_payload.get("_id")
if not isinstance(doc_id, str):
continue
status = op_payload.get("status")
# Deleting a missing doc is not a real failure -- the row is
# already gone, so we don't carry it forward on every flush.
if op_name == "delete" and status == 404:
continue
if status is None or status in _RETRYABLE_BULK_STATUSES:
retryable.add(doc_id)
else:
non_retryable.append(
_FailedBulkOp(
op=op_name,
doc_id=doc_id,
status=status if isinstance(status, int) else None,
error=_summarize_bulk_error(op_payload.get("error")),
)
)
return retryable, non_retryable
# Flush-time bulk batching limits. opensearch-py's helpers.async_bulk already
# splits a request by payload-byte budget (primary) and record count
# (secondary) via _ActionChunker -- semantically identical to MongoDB's
# _chunk_by_budget. We only expose those two limiter dimensions as env vars and
# pass them through as `max_chunk_bytes` / `chunk_size`, mirroring the MONGO_*
# knobs (lightrag/kg/mongo_impl.py) so behaviour stays consistent across
# backends. Defaults are tuned for OpenSearch: 100 MiB sits at the typical
# `http.max_content_length` ceiling, while the record caps match Mongo's.
DEFAULT_OPENSEARCH_UPSERT_MAX_PAYLOAD_BYTES = 100 * 1024 * 1024 # 100 MiB
DEFAULT_OPENSEARCH_UPSERT_MAX_RECORDS_PER_BATCH = 128
DEFAULT_OPENSEARCH_DELETE_MAX_RECORDS_PER_BATCH = 1000
# Sentinel "effectively unbounded" byte budget when payload splitting is
# disabled (env value <= 0). async_bulk needs a positive int here, so we use a
# large finite value in place of Mongo's float("inf").
_OPENSEARCH_UNBOUNDED_PAYLOAD_BYTES = 1 << 62
def _resolve_bulk_batch_limits() -> tuple[int, int, int]:
"""Resolve flush-time bulk batching limits from env, with module defaults.
Shared by every OpenSearch write path so the byte/record caps that bound a
single ``async_bulk`` request are consistent across all of them. A
non-positive value disables that splitting dimension (see
``_run_chunked_async_bulk``). Returns
``(upsert_payload_bytes, upsert_records, delete_records)``.
"""
upsert_payload_bytes = int(
_get_opensearch_env(
"OPENSEARCH_UPSERT_MAX_PAYLOAD_BYTES",
str(DEFAULT_OPENSEARCH_UPSERT_MAX_PAYLOAD_BYTES),
)
)
upsert_records = int(
_get_opensearch_env(
"OPENSEARCH_UPSERT_MAX_RECORDS_PER_BATCH",
str(DEFAULT_OPENSEARCH_UPSERT_MAX_RECORDS_PER_BATCH),
)
)
delete_records = int(
_get_opensearch_env(
"OPENSEARCH_DELETE_MAX_RECORDS_PER_BATCH",
str(DEFAULT_OPENSEARCH_DELETE_MAX_RECORDS_PER_BATCH),
)
)
if upsert_payload_bytes <= 0:
logger.warning(
f"OPENSEARCH_UPSERT_MAX_PAYLOAD_BYTES={upsert_payload_bytes} is non-positive, disable payload-size splitting"
)
if upsert_records <= 0:
logger.warning(
f"OPENSEARCH_UPSERT_MAX_RECORDS_PER_BATCH={upsert_records} is non-positive, disable upsert record-count splitting"
)
if delete_records <= 0:
logger.warning(
f"OPENSEARCH_DELETE_MAX_RECORDS_PER_BATCH={delete_records} is non-positive, disable delete record-count splitting"
)
return upsert_payload_bytes, upsert_records, delete_records
async def _run_chunked_async_bulk(
client: Any,
actions: list[dict[str, Any]],
*,
max_payload_bytes: int,
max_records_per_batch: int,
log_prefix: str,
what: str,
raise_on_error: bool = False,
**bulk_kwargs: Any,
) -> tuple[int, list[Any]]:
"""Run ``helpers.async_bulk`` with payload-size/record-count bounded chunks.
A thin wrapper that mirrors ``mongo_impl._run_batched_bulk_write`` in shape,
but delegates the actual splitting to opensearch-py's ``_ActionChunker``
(byte budget primary, record count secondary, oversized single action
emitted as its own chunk -- the same semantics as Mongo's
``_chunk_by_budget``). A non-positive limit disables that dimension. Extra
keyword arguments (e.g. ``refresh``) are forwarded to ``async_bulk``.
Returns ``async_bulk``'s ``(success, failed)`` tuple (``failed`` is empty
when ``raise_on_error=True``).
"""
if not actions:
return 0, []
chunk_size = max_records_per_batch if max_records_per_batch > 0 else len(actions)
max_chunk_bytes = (
max_payload_bytes
if max_payload_bytes > 0
else _OPENSEARCH_UNBOUNDED_PAYLOAD_BYTES
)
if len(actions) > chunk_size:
# Log format aligned with mongo_impl's flush split log
# (max_payload=/batch= field names, raw configured values). Unlike
# Mongo we cannot report the final batch count up front: async_bulk's
# _ActionChunker decides it at stream time by byte budget, so this
# record-count condition only catches count-driven splits.
logger.info(
f"{log_prefix} {what} split for {len(actions)} records "
f"(max_payload={max_payload_bytes} batch={max_records_per_batch})"
)
return await helpers.async_bulk(
client,
actions,
chunk_size=chunk_size,
max_chunk_bytes=max_chunk_bytes,
raise_on_error=raise_on_error,
**bulk_kwargs,
)
# Index _meta flag marking that an edges index has been migrated to canonical
# (sorted-pair) document ids. Guards the one-time reindex in
# PGGraphStorage-style startup so it runs at most once per index.
_EDGE_ID_CANONICAL_META_FLAG = "edge_id_canonical_v1"
# Emit a migration progress line every this many scanned edges, so operators
# watching a large-index reindex see liveness and an X/total denominator.
_EDGE_MIGRATION_PROGRESS_INTERVAL = 50_000
def _canonical_edge_id(source_node_id: str, target_node_id: str) -> str:
"""Direction-independent edge document ``_id``.
``hash(sorted(src, tgt))`` collapses an edge and its reverse onto the same
``_id``, so concurrent ``(A,B)``/``(B,A)`` writes overwrite one document
(last-write-wins) instead of racing into two separate docs. This makes
``upsert_edge`` idempotent by construction — no ``exists(reverse)``
read-then-write and no lock needed. The canonical id is always one of the
two directed ids ``hash("src-tgt")``/``hash("tgt-src")``, so the
bidirectional ``mget`` in ``has_edge``/``get_edge`` keeps finding it.
"""
lo, hi = sorted((source_node_id, target_node_id))
return compute_mdhash_id(f"{lo}-{hi}", prefix="edge-")
def _edge_source_id_list(doc: dict[str, Any]) -> list[str]:
"""Return an edge doc's source ids, from the ``source_ids`` array or by
splitting the ``GRAPH_FIELD_SEP``-joined ``source_id`` string."""
sids = doc.get("source_ids")
if not sids and doc.get("source_id"):
sids = doc["source_id"].split(GRAPH_FIELD_SEP)
return list(sids or [])
def _coerce_weight(weight: Any) -> float | None:
"""Coerce a (possibly string) edge weight to float, or None if non-numeric."""
if weight is None:
return None
try:
return float(weight)
except (TypeError, ValueError):
return None
def _merge_edge_payloads(docs: list[dict[str, Any]]) -> dict[str, Any]:
"""Merge edge-doc relation payloads when consolidating legacy duplicates.
``docs[0]`` is the survivor/base; the rest are duplicates folded into it.
Mirrors ``mongo_impl``'s dedupe merge and ``operate.py``'s
``_merge_edges_then_upsert`` field semantics (minus LLM description
summarisation): ``source_id``/``source_ids``/``file_path``/``description``
union their ``GRAPH_FIELD_SEP`` components, ``keywords`` are comma-set-
unioned, and ``weight`` is **summed across every fragment** (base + each
duplicate). Returns only the merged fields (to be layered onto the surviving
doc).
Weight summing deliberately does NOT dedup by ``source_id``: just like
``_merge_edges_then_upsert``, every edge fragment contributes its weight even
when fragments share a source/chunk id — reciprocal duplicates that came from
the same chunk still carry separate accumulated weight, so skipping them would
undercount the relation. This function is therefore not idempotent on its own;
idempotency across fail-fast retries is a property of the migration flow, not
this math: each folded reverse doc is deleted right after the canonical write
(see ``_merge_into_canonical_edge``), so a re-scan never re-presents an
already-folded reverse for summing. Legacy string weights are coerced;
non-numeric values are skipped so a bad value cannot crash the migration.
"""
source_ids: list[str] = []
file_paths: list[str] = []
descriptions: list[str] = []
keywords: set[str] = set()
weights: list[float] = []
for d in docs:
source_ids = merge_source_ids(source_ids, _edge_source_id_list(d))
fp = d.get("file_path")
file_paths = merge_source_ids(
file_paths, fp.split(GRAPH_FIELD_SEP) if fp else []
)
desc = d.get("description")
descriptions = merge_source_ids(
descriptions, desc.split(GRAPH_FIELD_SEP) if desc else []
)
kw = d.get("keywords")
if kw:
keywords.update(k.strip() for k in kw.split(",") if k.strip())
dw = _coerce_weight(d.get("weight"))
if dw is not None:
weights.append(dw)
merged: dict[str, Any] = {}
if source_ids:
merged["source_ids"] = source_ids
merged["source_id"] = GRAPH_FIELD_SEP.join(source_ids)
if file_paths:
merged["file_path"] = GRAPH_FIELD_SEP.join(file_paths)
if descriptions:
merged["description"] = GRAPH_FIELD_SEP.join(descriptions)
if keywords:
merged["keywords"] = ",".join(sorted(keywords))
if weights:
merged["weight"] = sum(weights)
return merged
# Detected at first connection; True when OpenSearch >= 3.3.0.
_shard_doc_supported: bool | None = None
def _pit_sort_with_field(field: str) -> list[dict]:
"""Return PIT sort clause with a unique field as primary sort.
Used purely as a pagination tiebreaker — order is fixed to asc since the
business sort (when present) is applied separately by the caller.
>= 3.3.0: _shard_doc only (most efficient, already unique within PIT).
< 3.3.0: field + _doc (field is unique, _doc for efficiency).
"""
if _shard_doc_supported:
return [{"_shard_doc": "asc"}]
return [{field: {"order": "asc"}}, {"_doc": "asc"}]
def _pit_sort_with_composite_key(*fields: str) -> list[dict]:
"""Return PIT sort clause with multiple fields forming a composite unique key.
>= 3.3.0: _shard_doc (most efficient, ignores the fields).
< 3.3.0: field1 + field2 + ... + _doc (composite is unique, _doc for efficiency).
"""
if _shard_doc_supported:
return [{"_shard_doc": "asc"}]
return [{f: {"order": "asc"}} for f in fields] + [{"_doc": "asc"}]
async def _detect_shard_doc_support(client: AsyncOpenSearch) -> bool:
"""Check if the cluster supports _shard_doc (OpenSearch >= 3.3.0)."""
try:
info = await client.info()
version_str = info.get("version", {}).get("number", "0.0.0")
# Strip pre-release suffixes (e.g. "3.3.0-SNAPSHOT" → "3", "3", "0")
parts = [p.split("-")[0] for p in version_str.split(".")]
major = int(parts[0]) if parts[0].isdigit() else 0
minor = int(parts[1]) if len(parts) > 1 and parts[1].isdigit() else 0
supported = (major > 3) or (major == 3 and minor >= 3)
logger.info(
f"OpenSearch version {version_str}: "
f"_shard_doc {'supported' if supported else 'not supported, using field+_doc fallback'}"
)
return supported
except Exception as e:
logger.warning(
f"Failed to detect OpenSearch version, assuming _shard_doc not supported: {e}"
)
return False
class ClientManager:
"""Singleton manager for OpenSearch client connections."""
_instances = {"client": None, "ref_count": 0}
_lock = asyncio.Lock()
@classmethod
async def get_client(cls) -> AsyncOpenSearch:
"""Get or create a shared AsyncOpenSearch client with reference counting."""
global _shard_doc_supported
async with cls._lock:
if cls._instances["client"] is None:
hosts_str = _get_opensearch_env("OPENSEARCH_HOSTS", "localhost:9200")
hosts = [h.strip() for h in hosts_str.split(",") if h.strip()]
username = _get_opensearch_env("OPENSEARCH_USER", "admin")
password = _get_opensearch_env("OPENSEARCH_PASSWORD", "admin")
use_ssl = _get_opensearch_env("OPENSEARCH_USE_SSL", "true").lower() in (
"true",
"1",
"yes",
)
verify_certs = _get_opensearch_env(
"OPENSEARCH_VERIFY_CERTS", "false"
).lower() in ("true", "1", "yes")
timeout = int(_get_opensearch_env("OPENSEARCH_TIMEOUT", "30"))
max_retries = int(_get_opensearch_env("OPENSEARCH_MAX_RETRIES", "3"))
ssl_context = None
if use_ssl and not verify_certs:
ssl_context = ssl_module.create_default_context()
ssl_context.check_hostname = False
ssl_context.verify_mode = ssl_module.CERT_NONE
client = AsyncOpenSearch(
hosts=hosts,
http_auth=(username, password) if username else None,
use_ssl=use_ssl,
verify_certs=verify_certs,
ssl_context=ssl_context,
ssl_show_warn=False,
timeout=timeout,
max_retries=max_retries,
retry_on_timeout=True,
)
cls._instances["client"] = client
cls._instances["ref_count"] = 0
_shard_doc_supported = await _detect_shard_doc_support(client)
logger.info(f"OpenSearch client connected to {hosts}")
cls._instances["ref_count"] += 1
return cls._instances["client"]
@classmethod
async def release_client(cls, client: AsyncOpenSearch):
"""Release a client reference. Closes the connection when ref count reaches 0."""
global _shard_doc_supported
async with cls._lock:
if client is not None and client is cls._instances["client"]:
cls._instances["ref_count"] -= 1
if cls._instances["ref_count"] <= 0:
try:
await cls._instances["client"].close()
except Exception:
pass
cls._instances["client"] = None
cls._instances["ref_count"] = 0
_shard_doc_supported = None
logger.info("OpenSearch client connection closed")
def _resolve_workspace(workspace: str, namespace: str):
"""Resolve effective workspace from env or parameter."""
opensearch_workspace = os.environ.get("OPENSEARCH_WORKSPACE")
if opensearch_workspace and opensearch_workspace.strip():
effective = opensearch_workspace.strip()
logger.info(
f"Using OPENSEARCH_WORKSPACE: '{effective}' (overriding '{workspace}/{namespace}')"
)
return effective
return workspace
def _build_index_name(workspace: str, namespace: str) -> tuple[str, str, str]:
"""Build index name and return (effective_workspace, final_namespace, index_name)."""
effective = _resolve_workspace(workspace, namespace)
if effective:
final_ns = f"{effective}_{namespace}"
else:
final_ns = namespace
effective = ""
index_name = _sanitize_index_name(final_ns)
return effective, final_ns, index_name
async def _mget_optional_doc(
client: AsyncOpenSearch,
index_name: str,
doc_id: str,
source_excludes: list[str] | None = None,
) -> dict[str, Any] | None:
"""Fetch a single document via mget and return None when it is absent.
``source_excludes`` is forwarded to OpenSearch's ``_source_excludes`` so
callers can ask the server to omit specific fields (e.g. ``["vector"]``)
and save network bandwidth.
"""
kwargs: dict[str, Any] = {"index": index_name, "body": {"ids": [doc_id]}}
if source_excludes:
kwargs["_source_excludes"] = source_excludes
response = await client.mget(**kwargs)
docs = response.get("docs", [])
if not docs:
return None
doc = docs[0]
if not doc.get("found"):
return None
return doc
def _is_missing_index_error(exc: Exception) -> bool:
"""Return True when an OpenSearch exception means the target index is missing."""
return "index_not_found_exception" in str(exc)
async def _verify_mirrored_id_mapping(client: AsyncOpenSearch, index_name: str) -> None:
"""Fail-fast when an existing index lacks the __mirrored_id keyword mapping.
Only enforced on OpenSearch < 3.3.0, where __mirrored_id serves as the
cross-shard pagination tiebreaker. Indices created by older LightRAG
releases will be missing this mapping; sorting by a missing field on a
multi-shard index can drop or duplicate documents during PIT pagination.
"""
if _shard_doc_supported:
return
try:
mapping = await client.indices.get_mapping(index=index_name)
except OpenSearchException:
return
props = mapping.get(index_name, {}).get("mappings", {}).get("properties", {})
if "__mirrored_id" not in props:
raise RuntimeError(
f"Index '{index_name}' lacks the '__mirrored_id' keyword mapping "
f"required for stable PIT pagination on OpenSearch < 3.3.0. "
f"This index was likely created by an older LightRAG release. "
f"Please reindex the data, or upgrade the cluster to OpenSearch >= 3.3.0."
)
@final
@dataclass
class OpenSearchKVStorage(BaseKVStorage):
"""Key-Value storage using OpenSearch. Uses dynamic mapping to support varied schemas."""
client: AsyncOpenSearch = field(default=None)
_index_name: str = field(default="", init=False)
_index_ready: bool = field(default=False, init=False)
def __init__(self, namespace, global_config, embedding_func, workspace=None):
super().__init__(
namespace=namespace,
workspace=workspace or "",
global_config=global_config,
embedding_func=embedding_func,
)
self.__post_init__()
def __post_init__(self):
validate_workspace(self.workspace)
self.workspace, self.final_namespace, self._index_name = _build_index_name(
self.workspace, self.namespace
)
# Pending writes are flushed via _flush_pending_kv_ops() during
# index_done_callback() / finalize(). Buffering many small upsert()
# invocations into a single async_bulk roundtrip avoids the per-call
# HTTP overhead profiled in issue #2785; the lock-everywhere model
# mirrors what #3043 introduced for OpenSearchVectorDBStorage.
self._pending_upserts: dict[str, dict[str, Any]] = {}
self._pending_kv_deletes: set[str] = set()
# Namespace-keyed lock (multi-process aware) is assigned in
# initialize(). All buffer reads / writes and the flush itself
# acquire this lock so an in-flight flush cannot interleave with
# concurrent get_by_id / upsert / delete on the same workspace.
self._flush_lock = None
(
self._max_upsert_payload_bytes,
self._max_upsert_records_per_batch,
self._max_delete_records_per_batch,
) = _resolve_bulk_batch_limits()
async def initialize(self):
"""Initialize client connection and create index if needed."""
async with get_data_init_lock():
if self.client is None:
self.client = await ClientManager.get_client()
await self._create_index_if_not_exists()
self._index_ready = True
logger.debug(
f"[{self.workspace}] OpenSearch KV storage initialized: {self._index_name}"
)
if self._flush_lock is None:
self._flush_lock = get_namespace_lock(
self.namespace, workspace=self.workspace
)
async def _ensure_index_ready(self):
"""Recreate the KV index after drop before the next write."""
if self._index_ready:
return
async with get_data_init_lock():
if self.client is None:
self.client = await ClientManager.get_client()
if not self._index_ready:
await self._create_index_if_not_exists()
self._index_ready = True
def _mark_index_missing(self):
"""Mark the KV index as unavailable for subsequent read short-circuiting."""
self._index_ready = False
async def _create_index_if_not_exists(self):
try:
if not await self.client.indices.exists(index=self._index_name):
# Use dynamic mapping so any namespace schema works
body = {
"mappings": {
"dynamic": True,
"properties": {
"__mirrored_id": {"type": "keyword"},
},
},
"settings": {
"index": {
"number_of_shards": _get_index_number_of_shards(),
"number_of_replicas": _get_index_number_of_replicas(),
},
},
}
await self.client.indices.create(index=self._index_name, body=body)
logger.info(f"[{self.workspace}] Created index: {self._index_name}")
else:
await _verify_mirrored_id_mapping(self.client, self._index_name)
except RequestError as e:
if "resource_already_exists_exception" not in str(e):
raise
except OpenSearchException as e:
logger.error(f"[{self.workspace}] Error creating index: {e}")
raise
async def finalize(self):
"""Flush pending writes and release the OpenSearch client connection.
Regular flush failures (any ``Exception``) are captured so they
can be re-surfaced as a ``RuntimeError`` that names the unflushed
buffer counts -- otherwise ``LightRAG.finalize_storages()`` would
log the storage as successfully finalized while writes silently
failed to reach OpenSearch.
``BaseException`` subclasses other than ``Exception`` (notably
``asyncio.CancelledError`` / ``KeyboardInterrupt`` / ``SystemExit``)
are NOT caught: they propagate through the ``finally`` block so
shutdown cancellation is honoured and not silently swallowed.
The client is released in ``finally`` so it does not leak whether
the flush succeeded, failed, or was cancelled.
"""
flush_error: Exception | None = None
try:
try:
await self._flush_pending_kv_ops()
except Exception as e:
# _flush_pending_kv_ops leaves the buffers intact on raise.
flush_error = e
finally:
if self.client is not None:
await ClientManager.release_client(self.client)
self.client = None
# Reached only when no BaseException propagated through the
# finally above. Snapshot remaining buffer state to report
# concrete counts.
pending_upserts = len(self._pending_upserts)
pending_deletes = len(self._pending_kv_deletes)
if flush_error is not None:
raise RuntimeError(
f"[{self.workspace}] OpenSearchKVStorage.finalize() flush "
f"raised; {pending_upserts} pending upserts and "
f"{pending_deletes} pending deletes were left buffered "
f"(client released, data lost)"
) from flush_error
if pending_upserts or pending_deletes:
raise RuntimeError(
f"[{self.workspace}] OpenSearchKVStorage.finalize() left "
f"{pending_upserts} pending upserts and {pending_deletes} "
f"pending deletes buffered after final flush attempt "
f"(transient bulk failure); these writes have been lost"
)
async def _iter_raw_docs(
self, batch_size: int = 1000
) -> AsyncIterator[list[dict[str, Any]]]:
"""Yield raw OpenSearch hits using PIT + search_after pagination."""
if not self._index_ready:
return
try:
pit = await self.client.create_pit(
index=self._index_name, params={"keep_alive": "1m"}
)
pit_id = pit["pit_id"]
try:
search_after = None
while True:
body = {
"query": {"match_all": {}},
"size": batch_size,
"pit": {"id": pit_id, "keep_alive": "1m"},
"sort": _pit_sort_with_field("__mirrored_id"),
}
if search_after:
body["search_after"] = search_after
response = await self.client.search(body=body)
hits = response["hits"]["hits"]
if not hits:
break
yield hits
search_after = hits[-1]["sort"]
if len(hits) < batch_size:
break
finally:
try:
await self.client.delete_pit(body={"pit_id": [pit_id]})
except Exception:
pass
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return
logger.error(f"[{self.workspace}] Error scanning documents: {e}")
raise
def _materialize_pending_kv_doc(
self, doc_id: str, source: dict[str, Any]
) -> dict[str, Any]:
"""Return a get_by_id-shaped view of a buffered upsert.
Mirrors the post-processing applied to mget hits: drops the
``__mirrored_id`` PIT sort key, attaches the ``_id`` field and
ensures ``create_time`` / ``update_time`` defaults are populated.
The buffer entry itself is not mutated.
"""
doc = {k: v for k, v in source.items() if k != "__mirrored_id"}
doc["_id"] = doc_id
doc.setdefault("create_time", 0)
doc.setdefault("update_time", 0)
return doc
async def get_by_id(self, id: str) -> dict[str, Any] | None:
"""Get a document by its ID, with read-your-writes against the buffer.
Priority: pending delete (tombstone) → pending upsert (buffered
write) → OpenSearch via mget. The buffered path strips
``__mirrored_id`` so the returned dict has the same shape as the
mget path.
"""
async with self._flush_lock:
if id in self._pending_kv_deletes:
return None
pending = self._pending_upserts.get(id)
if pending is not None:
return self._materialize_pending_kv_doc(id, pending)
if not self._index_ready:
return None
try:
response = await _mget_optional_doc(self.client, self._index_name, id)
if response is None:
return None
doc = response["_source"]
doc.pop("__mirrored_id", None)
doc["_id"] = response["_id"]
doc.setdefault("create_time", 0)
doc.setdefault("update_time", 0)
return doc
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return None
logger.error(f"[{self.workspace}] Error getting document {id}: {e}")
return None
async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
"""Get multiple documents by IDs (read-your-writes), preserving order.
Buffer is consulted under the lock with the same three-tier
priority as ``get_by_id``; remaining ids fall through to mget
outside the lock so the network call does not stall the flush.
"""
if not ids:
return []
buffered: dict[str, dict[str, Any] | None] = {}
remaining: list[str] = []
async with self._flush_lock:
for doc_id in ids:
if doc_id in self._pending_kv_deletes:
buffered[doc_id] = None
continue
pending = self._pending_upserts.get(doc_id)
if pending is not None:
buffered[doc_id] = self._materialize_pending_kv_doc(doc_id, pending)
continue
remaining.append(doc_id)
index_ready = self._index_ready
doc_map: dict[str, dict[str, Any] | None] = {}
if remaining and index_ready:
try:
response = await self.client.mget(
index=self._index_name, body={"ids": remaining}
)
for doc in response["docs"]:
if doc.get("found"):
data = doc["_source"]
data.pop("__mirrored_id", None)
data["_id"] = doc["_id"]
data.setdefault("create_time", 0)
data.setdefault("update_time", 0)
doc_map[doc["_id"]] = data
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
else:
logger.error(f"[{self.workspace}] Error getting documents: {e}")
return [
buffered[doc_id] if doc_id in buffered else doc_map.get(doc_id)
for doc_id in ids
]
async def filter_keys(self, keys: set[str]) -> set[str]:
"""Return the subset of keys that do not exist in storage.
Buffer-aware: buffered upserts count as "exists" (and so are
removed from the missing set), buffered deletes count as
"missing" and are NOT queried via mget (a persisted-but-pending-
delete row would otherwise be misclassified as existing).
"""
async with self._flush_lock:
pending_upserts = set(self._pending_upserts)
pending_deletes = set(self._pending_kv_deletes)
index_ready = self._index_ready
# Buffered upserts shadow OpenSearch -- they will exist after flush.
to_check = keys - pending_upserts - pending_deletes
if not to_check:
# All keys are accounted for by the buffer alone.
return keys - pending_upserts
if not index_ready:
return keys - pending_upserts
try:
response = await self.client.mget(
index=self._index_name,
body={"ids": list(to_check)},
_source=False,
)
existing_on_server = {
doc["_id"] for doc in response["docs"] if doc.get("found")
}
return (keys - pending_upserts) - existing_on_server
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return keys - pending_upserts
logger.error(f"[{self.workspace}] Error filtering keys: {e}")
return keys - pending_upserts
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
"""Buffer documents for batched flush.
Time-stamping and ``__mirrored_id`` injection happen eagerly so the
persisted shape matches what reads expect; the actual ``async_bulk``
call is deferred to ``_flush_pending_kv_ops()`` invoked from
``index_done_callback`` / ``finalize``.
Multi-worker note: the buffer is process-local. Other workers will
not see these writes until ``index_done_callback()`` flushes them.
"""
if not data:
return
await self._ensure_index_ready()
logger.debug(
f"[{self.workspace}] Buffering {len(data)} documents for {self.namespace}"
)
current_time = int(time.time())
# Construct sources outside the lock (no IO; just dict shuffling)
# so we hold the lock only for the buffer-swap step.
prepared: list[tuple[str, dict[str, Any]]] = []
for i, (doc_id, doc_data) in enumerate(data.items(), start=1):
doc_data["update_time"] = current_time
doc_data.setdefault("create_time", current_time)
source = {k: v for k, v in doc_data.items() if k != "_id"}
source["__mirrored_id"] = doc_id
prepared.append((doc_id, source))
await _cooperative_yield(i)
# Buffer: an upsert cancels any pending delete on the same id.
async with self._flush_lock:
for doc_id, source in prepared:
self._pending_kv_deletes.discard(doc_id)
self._pending_upserts[doc_id] = source
async def delete(self, ids: list[str]) -> None:
"""Buffer document deletes for batched flush.
A delete cancels any pending upsert on the same id; the actual
bulk delete is performed by ``_flush_pending_kv_ops`` during the
next ``index_done_callback`` / ``finalize`` call.
``_index_ready`` is intentionally NOT checked here: even if the
index has been marked missing, the buffered upsert (if any) must
still be invalidated, otherwise a subsequent flush would resurrect
a logically-deleted key.
"""
if not ids:
return
if isinstance(ids, set):
ids = list(ids)
async with self._flush_lock:
for doc_id in ids:
self._pending_upserts.pop(doc_id, None)
self._pending_kv_deletes.add(doc_id)
logger.debug(
f"[{self.workspace}] Buffered delete for {len(ids)} documents in {self.namespace}"
)
async def _flush_pending_kv_ops(self) -> None:
"""Flush buffered upserts + deletes via a single async_bulk call.
Concurrency contract: the entire flush runs under ``_flush_lock``;
``upsert`` / ``delete`` / reads / ``drop`` all acquire the same lock
so an in-flight flush cannot interleave with concurrent buffer
mutations.
Failure handling mirrors the Vector-side helper:
* If ``_ensure_index_ready`` raises, the buffers are left intact
and the next flush retries.
* If ``async_bulk`` raises, the buffers are left intact.
* Per-doc retryable failures (408 / 429 / 5xx) stay in the buffer.
* Per-doc non-retryable failures (most 4xx) are cleared and a
sample is logged at WARNING with op / id / status / error.
"""
async with self._flush_lock:
if not self._pending_upserts and not self._pending_kv_deletes:
return
if self.client is None:
return
await self._ensure_index_ready()
pending_upserts = self._pending_upserts
pending_deletes = self._pending_kv_deletes
# Deletes are flushed before upserts so a delete followed (in time)
# by an upsert on the same id still ends as an index; the two
# buffers are disjoint anyway (upsert/delete pop each other), so
# running them as separate async_bulk requests is safe and lets the
# delete record-count cap differ from the upsert cap (mirrors
# mongo_impl's separate upsert/delete phases).
delete_actions: list[dict[str, Any]] = [
{
"_op_type": "delete",
"_index": self._index_name,
"_id": doc_id,
}
for doc_id in pending_deletes
]
index_actions: list[dict[str, Any]] = [
{
"_op_type": "index",
"_index": self._index_name,
"_id": doc_id,
"_source": source,
}
for doc_id, source in pending_upserts.items()
]
try:
log_prefix = f"[{self.workspace}] {self.namespace} flush:"
del_success, del_failed = await _run_chunked_async_bulk(
self.client,
delete_actions,
max_payload_bytes=self._max_upsert_payload_bytes,
max_records_per_batch=self._max_delete_records_per_batch,
log_prefix=log_prefix,
what="delete",
raise_on_error=False,
)
idx_success, idx_failed = await _run_chunked_async_bulk(
self.client,
index_actions,
max_payload_bytes=self._max_upsert_payload_bytes,
max_records_per_batch=self._max_upsert_records_per_batch,
log_prefix=log_prefix,
what="upsert",
raise_on_error=False,
)
success = del_success + idx_success
failed = list(del_failed) + list(idx_failed)
except OpenSearchException as e:
logger.error(
f"[{self.workspace}] Error flushing KV ops "
f"(upserts={len(pending_upserts)}, "
f"deletes={len(pending_deletes)}): {e}"
)
raise
retryable_ids, non_retryable_ops = _extract_bulk_failed_ids(failed)
non_retryable_ids = {op.doc_id for op in non_retryable_ops}
# Keep ONLY retryable ops buffered for the next flush. Successful
# ops are popped; non-retryable (permanent 4xx) ops are dropped
# here, not retained: a permanently-unwritable op can never land,
# so keeping it would replay-and-refail on every later flush and
# poison every caller that shares this buffer — including direct
# flush paths (e.g. _persist_parsed_full_docs) that never run the
# pipeline's cleanup. The raise below (not retention) is what
# surfaces the failure and prevents a silent PROCESSED.
keep_ids = retryable_ids
for doc_id in list(pending_upserts.keys()):
if doc_id not in keep_ids:
pending_upserts.pop(doc_id, None)
new_deletes: set[str] = {
doc_id for doc_id in pending_deletes if doc_id in keep_ids
}
pending_deletes.clear()
pending_deletes.update(new_deletes)
if retryable_ids:
logger.warning(
f"[{self.workspace}] {len(retryable_ids)} KV ops will "
f"retry on the next flush (transient failure)"
)
logger.debug(
f"[{self.workspace}] Flushed KV ops: {success} ok, "
f"retry={len(retryable_ids)}, permanent_fail={len(non_retryable_ids)}"
)
if non_retryable_ops:
sample = non_retryable_ops[:5]
sample_text = ", ".join(
f"{op.op}/{op.doc_id}/status={op.status}/{op.error}"
for op in sample
)
# A permanent (non-retryable) bulk failure means the data did
# not land. Raise so index_done_callback surfaces it and the
# pipeline aborts instead of marking the document PROCESSED.
raise RuntimeError(
f"[{self.workspace}] {self.namespace} flush: "
f"{len(non_retryable_ops)} KV ops failed permanently "
f"(non-retryable). Sample: {sample_text}"
)
async def drop_pending_index_ops(self) -> None:
"""Discard buffered upserts/deletes (pipeline aborting on error)."""
async with self._flush_lock:
self._pending_upserts.clear()
self._pending_kv_deletes.clear()
async def index_done_callback(self) -> None:
"""Flush pending KV ops and refresh the index for search visibility.
Flush runs first so a previously-missing index gets recreated by
``_flush_pending_kv_ops`` (via ``_ensure_index_ready``) before any
buffered writes are abandoned. The refresh step is skipped only
when the index is still not ready after the flush attempt.
"""
await self._flush_pending_kv_ops()
if not self._index_ready:
return
try:
await self.client.indices.refresh(index=self._index_name)
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return
raise
async def is_empty(self) -> bool:
"""Return True if the index (plus pending buffer) contains no docs.
Buffer-aware: a pending upsert makes is_empty False immediately,
avoiding the counterintuitive "I just upserted but is_empty
returned True" case. Pending deletes alone are not enough to flip
the answer because we cannot tell whether other persisted rows
survive without flushing.
"""
async with self._flush_lock:
if self._pending_upserts:
return False
index_ready = self._index_ready
if not index_ready:
return True
try:
response = await self.client.count(index=self._index_name)
return response["count"] == 0
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return True
async def drop(self) -> dict[str, str]:
"""Delete the entire index, discarding pending buffers.
Runs entirely under ``_flush_lock`` so a concurrent flush / upsert
cannot land writes against an index that is being deleted.
"""
async with self._flush_lock:
# Pending writes are meaningless once the index is dropped.
self._pending_upserts.clear()
self._pending_kv_deletes.clear()
try:
try:
await self.client.indices.delete(index=self._index_name)
logger.info(f"[{self.workspace}] Dropped index: {self._index_name}")
except NotFoundError:
logger.info(
f"[{self.workspace}] Index already missing during drop: {self._index_name}"
)
self._mark_index_missing()
return {
"status": "success",
"message": f"Index {self._index_name} dropped",
}
except OpenSearchException as e:
self._mark_index_missing()
logger.error(f"[{self.workspace}] Error dropping index: {e}")
return {"status": "error", "message": str(e)}
except Exception as e:
self._mark_index_missing()
logger.error(f"[{self.workspace}] Unexpected error dropping index: {e}")
return {"status": "error", "message": str(e)}
@final
@dataclass
class OpenSearchDocStatusStorage(DocStatusStorage):
"""Document status storage using OpenSearch."""
client: AsyncOpenSearch = field(default=None)
_index_name: str = field(default="", init=False)
_index_ready: bool = field(default=False, init=False)
def __init__(self, namespace, global_config, embedding_func, workspace=None):
super().__init__(
namespace=namespace,
workspace=workspace or "",
global_config=global_config,
embedding_func=embedding_func,
)
self.__post_init__()
def __post_init__(self):
validate_workspace(self.workspace)
self.workspace, self.final_namespace, self._index_name = _build_index_name(
self.workspace, self.namespace
)
(
self._max_upsert_payload_bytes,
self._max_upsert_records_per_batch,
self._max_delete_records_per_batch,
) = _resolve_bulk_batch_limits()
def _prepare_doc_status_data(self, doc: dict[str, Any]) -> dict[str, Any]:
"""Normalize a raw OpenSearch document to DocProcessingStatus-compatible dict."""
data = doc.copy()
data.pop("_id", None)
data.pop("__mirrored_id", None)
if "file_path" not in data:
data["file_path"] = "no-file-path"
data.setdefault("metadata", {})
data.setdefault("error_msg", None)
if "error" in data:
if not data.get("error_msg"):
data["error_msg"] = data.pop("error")
else:
data.pop("error", None)
return data
async def initialize(self):
"""Initialize client connection and create doc status index."""
async with get_data_init_lock():
if self.client is None:
self.client = await ClientManager.get_client()
await self._create_index_if_not_exists()
self._index_ready = True
logger.debug(
f"[{self.workspace}] OpenSearch DocStatus storage initialized: {self._index_name}"
)
async def _ensure_index_ready(self):
"""Recreate the doc status index after drop before the next write."""
if self._index_ready:
return
async with get_data_init_lock():
if self.client is None:
self.client = await ClientManager.get_client()
if not self._index_ready:
await self._create_index_if_not_exists()
self._index_ready = True
def _mark_index_missing(self):
"""Mark the doc status index as unavailable for subsequent read short-circuiting."""
self._index_ready = False
async def _create_index_if_not_exists(self):
try:
if not await self.client.indices.exists(index=self._index_name):
body = {
"mappings": {
"dynamic": True,
"properties": {
"__mirrored_id": {"type": "keyword"},
"status": {"type": "keyword"},
"file_path": {"type": "keyword"},
"track_id": {"type": "keyword"},
"content_hash": {"type": "keyword"},
"created_at": {"type": "date"},
"updated_at": {"type": "date"},
},
},
"settings": {
"index": {
"number_of_shards": _get_index_number_of_shards(),
"number_of_replicas": _get_index_number_of_replicas(),
},
},
}
await self.client.indices.create(index=self._index_name, body=body)
logger.info(
f"[{self.workspace}] Created doc status index: {self._index_name}"
)
else:
await _verify_mirrored_id_mapping(self.client, self._index_name)
await self._ensure_content_hash_mapping()
except RequestError as e:
if "resource_already_exists_exception" not in str(e):
raise
except OpenSearchException as e:
logger.error(f"[{self.workspace}] Error creating doc status index: {e}")
raise
async def _ensure_content_hash_mapping(self) -> None:
"""Add the content_hash keyword mapping to a pre-existing doc status index.
Indices created by older LightRAG releases lack content_hash entirely.
put_mapping is idempotent for new fields, so this is safe to call every
startup; we only fail loudly when the cluster reports a mapping conflict
(which would indicate dynamic mapping already coerced content_hash to a
different type).
"""
try:
mapping = await self.client.indices.get_mapping(index=self._index_name)
except OpenSearchException:
return
props = (
mapping.get(self._index_name, {}).get("mappings", {}).get("properties", {})
)
if "content_hash" in props:
return
try:
await self.client.indices.put_mapping(
index=self._index_name,
body={"properties": {"content_hash": {"type": "keyword"}}},
)
logger.info(
f"[{self.workspace}] Added content_hash keyword mapping to {self._index_name}"
)
except OpenSearchException as e:
logger.warning(
f"[{self.workspace}] Failed to add content_hash mapping to "
f"{self._index_name}: {e}"
)
async def finalize(self):
"""Release the OpenSearch client connection."""
if self.client is not None:
await ClientManager.release_client(self.client)
self.client = None
async def get_by_id(self, id: str) -> Union[dict[str, Any], None]:
"""Get a document status record by ID."""
if not self._index_ready:
return None
try:
response = await _mget_optional_doc(self.client, self._index_name, id)
if response is None:
return None
doc = response["_source"]
doc["_id"] = response["_id"]
return doc
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return None
logger.error(f"[{self.workspace}] Error getting doc status {id}: {e}")
return None
async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
"""Get multiple document status records by IDs."""
if not self._index_ready:
return [None] * len(ids)
try:
response = await self.client.mget(index=self._index_name, body={"ids": ids})
doc_map = {}
for doc in response["docs"]:
if doc.get("found"):
data = doc["_source"]
data["_id"] = doc["_id"]
doc_map[doc["_id"]] = data
return [doc_map.get(id) for id in ids]
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return [None] * len(ids)
logger.error(f"[{self.workspace}] Error getting doc statuses: {e}")
return [None] * len(ids)
async def filter_keys(self, keys: set[str]) -> set[str]:
"""Return the subset of keys that do not exist in storage."""
if not self._index_ready:
return keys
try:
response = await self.client.mget(
index=self._index_name, body={"ids": list(keys)}, _source=False
)
existing_ids = {doc["_id"] for doc in response["docs"] if doc.get("found")}
return keys - existing_ids
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return keys
logger.error(f"[{self.workspace}] Error filtering keys: {e}")
return keys
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
"""Insert or update document status records."""
if not data:
return
await self._ensure_index_ready()
logger.debug(f"[{self.workspace}] Upserting {len(data)} doc statuses")
actions = []
for i, (k, v) in enumerate(data.items(), start=1):
v.setdefault("chunks_list", [])
source = {fk: fv for fk, fv in v.items() if fk != "_id"}
source["__mirrored_id"] = k
actions.append(
{
"_op_type": "index",
"_index": self._index_name,
"_id": k,
"_source": source,
}
)
await _cooperative_yield(i)
try:
# DocStatus needs refresh="wait_for" because get_docs_by_status
# (search-based) is called immediately after enqueue upserts.
await _run_chunked_async_bulk(
self.client,
actions,
max_payload_bytes=self._max_upsert_payload_bytes,
max_records_per_batch=self._max_upsert_records_per_batch,
log_prefix=f"[{self.workspace}] {self.namespace} upsert:",
what="doc-status upsert",
raise_on_error=False,
refresh="wait_for",
)
except OpenSearchException as e:
logger.error(f"[{self.workspace}] Error upserting doc statuses: {e}")
# Surface the failure instead of returning as if the status write
# succeeded — a silently-lost doc-status row corrupts pipeline
# bookkeeping (a doc could never be recorded FAILED/PROCESSED).
raise
async def get_status_counts(self) -> dict[str, int]:
"""Get document counts grouped by status."""
if not self._index_ready:
return {}
try:
body = {
"size": 0,
"aggs": {"status_counts": {"terms": {"field": "status", "size": 100}}},
}
response = await self.client.search(index=self._index_name, body=body)
return {
bucket["key"]: bucket["doc_count"]
for bucket in response["aggregations"]["status_counts"]["buckets"]
}
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return {}
logger.error(f"[{self.workspace}] Error getting status counts: {e}")
return {}
async def _search_all_docs(self, query: dict) -> dict[str, DocProcessingStatus]:
"""Fetch all documents matching a query using PIT + search_after."""
if not self._index_ready:
return {}
result = {}
batch_size = 10000
try:
pit = await self.client.create_pit(
index=self._index_name, params={"keep_alive": "1m"}
)
pit_id = pit["pit_id"]
try:
search_after = None
while True:
body = {
"query": query,
"size": batch_size,
"pit": {"id": pit_id, "keep_alive": "1m"},
"sort": _pit_sort_with_field("__mirrored_id"),
}
if search_after:
body["search_after"] = search_after
response = await self.client.search(body=body)
hits = response["hits"]["hits"]
if not hits:
break
for hit in hits:
try:
data = self._prepare_doc_status_data(hit["_source"])
result[hit["_id"]] = DocProcessingStatus(**data)
except (KeyError, TypeError) as e:
logger.error(
f"[{self.workspace}] Error parsing doc {hit['_id']}: {e}"
)
search_after = hits[-1]["sort"]
if len(hits) < batch_size:
break
finally:
try:
await self.client.delete_pit(body={"pit_id": [pit_id]})
except Exception:
pass
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return {}
logger.error(f"[{self.workspace}] Error fetching docs: {e}")
return result
async def get_docs_by_status(
self, status: DocStatus
) -> dict[str, DocProcessingStatus]:
"""Get all documents matching a specific processing status."""
return await self.get_docs_by_statuses([status])
async def get_docs_by_statuses(
self, statuses: list[DocStatus]
) -> dict[str, DocProcessingStatus]:
"""Get all documents matching any of the given statuses in a single query.
Uses OpenSearch's terms query (multi-value equivalent of term) to fetch
all matching statuses in one PIT + search_after pass instead of one
full scan per status.
"""
if not statuses:
return {}
status_values = [s.value for s in statuses]
return await self._search_all_docs({"terms": {"status": status_values}})
async def get_docs_by_track_id(
self, track_id: str
) -> dict[str, DocProcessingStatus]:
"""Get all documents matching a specific track ID."""
return await self._search_all_docs({"term": {"track_id": track_id}})
async def get_docs_paginated(
self,
status_filter: DocStatus | None = None,
status_filters: list[DocStatus] | None = None,
page: int = 1,
page_size: int = 50,
sort_field: str = "updated_at",
sort_direction: str = "desc",
) -> tuple[list[tuple[str, DocProcessingStatus]], int]:
"""Get documents with pagination using PIT + search_after."""
if not self._index_ready:
return [], 0
status_filter_values = self.resolve_status_filter_values(
status_filter=status_filter,
status_filters=status_filters,
)
page = max(1, page)
page_size = max(10, min(200, page_size))
if sort_field == "id":
sort_field = "_id"
if sort_field not in ("created_at", "updated_at", "_id", "file_path"):
sort_field = "updated_at"
sort_order = "asc" if sort_direction.lower() == "asc" else "desc"
query = {"match_all": {}}
if status_filter_values is not None:
if len(status_filter_values) == 1:
query = {"term": {"status": next(iter(status_filter_values))}}
else:
query = {"terms": {"status": sorted(status_filter_values)}}
skip_count = (page - 1) * page_size
try:
count_resp = await self.client.count(
index=self._index_name, body={"query": query}
)
total_count = count_resp.get("count", 0)
if total_count == 0 or skip_count >= total_count:
return [], total_count
sort_clause = [{sort_field: {"order": sort_order}}] + _pit_sort_with_field(
"__mirrored_id"
)
pit = await self.client.create_pit(
index=self._index_name, params={"keep_alive": "1m"}
)
pit_id = pit["pit_id"]
try:
search_after = None
skipped = 0
while skipped < skip_count:
batch = min(page_size, skip_count - skipped)
body = {
"query": query,
"sort": sort_clause,
"size": batch,
"pit": {"id": pit_id, "keep_alive": "1m"},
}
if search_after:
body["search_after"] = search_after
resp = await self.client.search(body=body)
hits = resp["hits"]["hits"]
if not hits:
return [], total_count
search_after = hits[-1]["sort"]
skipped += len(hits)
body = {
"query": query,
"sort": sort_clause,
"size": page_size,
"pit": {"id": pit_id, "keep_alive": "1m"},
}
if search_after:
body["search_after"] = search_after
response = await self.client.search(body=body)
finally:
try:
await self.client.delete_pit(body={"pit_id": [pit_id]})
except Exception:
pass
documents = []
for hit in response["hits"]["hits"]:
try:
data = self._prepare_doc_status_data(hit["_source"])
documents.append((hit["_id"], DocProcessingStatus(**data)))
except (KeyError, TypeError) as e:
logger.error(
f"[{self.workspace}] Error parsing doc {hit['_id']}: {e}"
)
return documents, total_count
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return [], 0
logger.error(f"[{self.workspace}] Error in paginated query: {e}")
return [], 0
async def get_all_status_counts(self) -> dict[str, int]:
"""Get document counts for all statuses including an 'all' total."""
if not self._index_ready:
return {}
try:
body = {
"size": 0,
"aggs": {"status_counts": {"terms": {"field": "status", "size": 100}}},
}
response = await self.client.search(index=self._index_name, body=body)
counts = {}
total = 0
for bucket in response["aggregations"]["status_counts"]["buckets"]:
counts[bucket["key"]] = bucket["doc_count"]
total += bucket["doc_count"]
counts["all"] = total
return counts
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return {}
logger.error(f"[{self.workspace}] Error getting all status counts: {e}")
return {}
async def get_doc_by_file_path(self, file_path: str) -> Union[dict[str, Any], None]:
"""Find a document status record by its file_path field."""
if not self._index_ready:
return None
try:
body = {"query": {"term": {"file_path": file_path}}, "size": 1}
response = await self.client.search(index=self._index_name, body=body)
hits = response["hits"]["hits"]
if hits:
doc = hits[0]["_source"]
doc["_id"] = hits[0]["_id"]
return doc
return None
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return None
logger.error(f"[{self.workspace}] Error getting doc by file_path: {e}")
return None
async def get_doc_by_file_basename(
self, basename: str
) -> Union[tuple[str, dict[str, Any]], None]:
"""Find an existing record whose canonical basename matches.
The caller is responsible for passing an already-canonical basename;
stored ``file_path`` values are canonicalized by the business layer, so
this lookup performs an exact term query against the file_path keyword
field.
"""
if not basename:
return None
if basename == "unknown_source":
return None
if not self._index_ready:
return None
try:
body = {"query": {"term": {"file_path": basename}}, "size": 1}
response = await self.client.search(index=self._index_name, body=body)
hits = response["hits"]["hits"]
if not hits:
return None
hit = hits[0]
doc = hit["_source"]
return hit["_id"], doc
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return None
logger.error(f"[{self.workspace}] Error getting doc by file_basename: {e}")
return None
async def get_doc_by_content_hash(
self, content_hash: str
) -> Union[tuple[str, dict[str, Any]], None]:
"""Find an existing record whose content_hash field matches.
Uses the content_hash keyword mapping created by
``_create_index_if_not_exists`` / ``_ensure_content_hash_mapping``.
Empty values short-circuit so legacy rows without the field cannot
accidentally match via type coercion.
"""
if not content_hash:
return None
if not self._index_ready:
return None
try:
body = {"query": {"term": {"content_hash": content_hash}}, "size": 1}
response = await self.client.search(index=self._index_name, body=body)
hits = response["hits"]["hits"]
if not hits:
return None
hit = hits[0]
doc = hit["_source"]
return hit["_id"], doc
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return None
logger.error(f"[{self.workspace}] Error getting doc by content_hash: {e}")
return None
async def index_done_callback(self) -> None:
"""Refresh index to make recently indexed documents searchable."""
if not self._index_ready:
return
try:
await self.client.indices.refresh(index=self._index_name)
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return
raise
async def is_empty(self) -> bool:
"""Return True if the index contains no documents."""
if not self._index_ready:
return True
try:
response = await self.client.count(index=self._index_name)
return response["count"] == 0
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return True
async def delete(self, ids: list[str]) -> None:
"""Delete document status records by IDs."""
if not ids:
return
if not self._index_ready:
return
if isinstance(ids, set):
ids = list(ids)
try:
# DocStatus needs refresh="wait_for" because downstream readers
# (get_docs_by_status, get_docs_paginated, etc.) are search-based
# and callers like _validate_and_fix_document_consistency() may
# query immediately after deletion without index_done_callback().
actions = [
{"_op_type": "delete", "_index": self._index_name, "_id": doc_id}
for doc_id in ids
]
await _run_chunked_async_bulk(
self.client,
actions,
max_payload_bytes=self._max_upsert_payload_bytes,
max_records_per_batch=self._max_delete_records_per_batch,
log_prefix=f"[{self.workspace}] {self.namespace} delete:",
what="doc-status delete",
raise_on_error=False,
refresh="wait_for",
)
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return
logger.error(f"[{self.workspace}] Error deleting doc statuses: {e}")
async def drop(self) -> dict[str, str]:
"""Delete the entire doc status index."""
try:
try:
await self.client.indices.delete(index=self._index_name)
logger.info(
f"[{self.workspace}] Dropped doc status index: {self._index_name}"
)
except NotFoundError:
logger.info(
f"[{self.workspace}] Doc status index already missing during drop: {self._index_name}"
)
self._mark_index_missing()
return {"status": "success", "message": f"Index {self._index_name} dropped"}
except OpenSearchException as e:
self._mark_index_missing()
logger.error(f"[{self.workspace}] Error dropping doc status index: {e}")
return {"status": "error", "message": str(e)}
except Exception as e:
self._mark_index_missing()
logger.error(
f"[{self.workspace}] Unexpected error dropping doc status index: {e}"
)
return {"status": "error", "message": str(e)}
@final
@dataclass
class OpenSearchGraphStorage(BaseGraphStorage):
"""Graph storage using OpenSearch with separate nodes and edges indices.
Supports two BFS traversal strategies:
- PPL graphlookup (server-side BFS, requires OpenSearch SQL plugin with Calcite engine)
- Application-level batched BFS (fallback, works on any OpenSearch 3.x+)
The strategy is auto-detected during initialize() and can be overridden via
the OPENSEARCH_USE_PPL_GRAPHLOOKUP environment variable (true/false).
"""
client: AsyncOpenSearch = field(default=None)
_nodes_index: str = field(default="", init=False)
_edges_index: str = field(default="", init=False)
_indices_ready: bool = field(default=False, init=False)
_nodes_dirty: bool = field(default=False, init=False)
_edges_dirty: bool = field(default=False, init=False)
_ppl_graphlookup_available: bool = field(default=False, init=False)
def __init__(self, namespace, global_config, embedding_func, workspace=None):
super().__init__(
namespace=namespace,
workspace=workspace or "",
global_config=global_config,
embedding_func=embedding_func,
)
self.__post_init__()
def __post_init__(self):
validate_workspace(self.workspace)
self.workspace, self.final_namespace, base_name = _build_index_name(
self.workspace, self.namespace
)
self._nodes_index = f"{base_name}-nodes"
self._edges_index = f"{base_name}-edges"
(
self._max_upsert_payload_bytes,
self._max_upsert_records_per_batch,
self._max_delete_records_per_batch,
) = _resolve_bulk_batch_limits()
async def initialize(self):
"""Initialize client, create indices, and detect PPL graphlookup support."""
async with get_data_init_lock():
if self.client is None:
self.client = await ClientManager.get_client()
await self._create_indices_if_not_exist()
await self._migrate_edges_to_canonical_id_if_needed()
self._indices_ready = True
self._nodes_dirty = False
self._edges_dirty = False
await self._detect_ppl_graphlookup()
logger.debug(
f"[{self.workspace}] OpenSearch Graph storage initialized: "
f"{self._nodes_index}, {self._edges_index} "
f"(PPL graphlookup: {self._ppl_graphlookup_available})"
)
async def _ensure_indices_ready(self):
"""Recreate graph indices after drop before the next write."""
if self._indices_ready:
return
async with get_data_init_lock():
if self.client is None:
self.client = await ClientManager.get_client()
if not self._indices_ready:
await self._create_indices_if_not_exist()
self._indices_ready = True
def _mark_indices_missing(self):
"""Mark graph indices as unavailable for subsequent read short-circuiting."""
self._indices_ready = False
self._nodes_dirty = False
self._edges_dirty = False
async def _refresh_graph_indices_if_dirty(
self, *, refresh_nodes: bool = False, refresh_edges: bool = False
) -> None:
"""Refresh graph indices only when prior writes made search views stale."""
if not self._indices_ready:
return
if not (
(refresh_nodes and self._nodes_dirty)
or (refresh_edges and self._edges_dirty)
):
return
try:
async with get_data_init_lock():
if refresh_nodes and self._nodes_dirty:
await self.client.indices.refresh(index=self._nodes_index)
self._nodes_dirty = False
if refresh_edges and self._edges_dirty:
await self.client.indices.refresh(index=self._edges_index)
self._edges_dirty = False
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return
raise
async def _detect_ppl_graphlookup(self):
"""Detect whether PPL graphlookup command is available on this cluster."""
env_override = os.environ.get("OPENSEARCH_USE_PPL_GRAPHLOOKUP", "").lower()
if env_override == "true":
self._ppl_graphlookup_available = True
return
if env_override == "false":
self._ppl_graphlookup_available = False
return
# Auto-detect by sending a minimal PPL query
try:
await self.client.transport.perform_request(
"POST",
"/_plugins/_ppl",
body={"query": f"source = {self._edges_index} | head 0"},
)
# PPL endpoint works; now test graphlookup syntax with a no-op query
await self.client.transport.perform_request(
"POST",
"/_plugins/_ppl",
body={
"query": (
f"source = {self._edges_index} | head 1 "
f"| graphLookup {self._edges_index} "
f"start=source_node_id edge=target_node_id-->source_node_id "
f"maxDepth=0 as _gl_probe"
)
},
)
self._ppl_graphlookup_available = True
logger.info(
f"[{self.workspace}] PPL graphlookup is available, using server-side BFS"
)
except Exception:
self._ppl_graphlookup_available = False
logger.info(
f"[{self.workspace}] PPL graphlookup not available, using client-side BFS"
)
async def _create_indices_if_not_exist(self):
try:
if not await self.client.indices.exists(index=self._nodes_index):
body = {
"mappings": {
"dynamic": True,
"properties": {
"entity_id": {"type": "keyword"},
"entity_type": {"type": "keyword"},
"description": {"type": "text"},
"source_id": {"type": "text"},
"source_ids": {"type": "keyword"},
"file_path": {"type": "keyword"},
"created_at": {"type": "long"},
},
},
"settings": {
"index": {
"number_of_shards": _get_index_number_of_shards(),
"number_of_replicas": _get_index_number_of_replicas(),
}
},
}
await self.client.indices.create(index=self._nodes_index, body=body)
logger.info(
f"[{self.workspace}] Created nodes index: {self._nodes_index}"
)
except RequestError as e:
if "resource_already_exists_exception" not in str(e):
raise
try:
if not await self.client.indices.exists(index=self._edges_index):
body = {
"mappings": {
"dynamic": True,
"properties": {
"source_node_id": {"type": "keyword"},
"target_node_id": {"type": "keyword"},
"relationship": {"type": "keyword"},
"description": {"type": "text"},
"weight": {"type": "float"},
"keywords": {"type": "text"},
"source_id": {"type": "text"},
"source_ids": {"type": "keyword"},
"file_path": {"type": "keyword"},
"created_at": {"type": "long"},
},
},
"settings": {
"index": {
"number_of_shards": _get_index_number_of_shards(),
"number_of_replicas": _get_index_number_of_replicas(),
}
},
}
await self.client.indices.create(index=self._edges_index, body=body)
logger.info(
f"[{self.workspace}] Created edges index: {self._edges_index}"
)
except RequestError as e:
if "resource_already_exists_exception" not in str(e):
raise
async def _migrate_edges_to_canonical_id_if_needed(self) -> None:
"""One-time reindex of edge docs onto canonical (sorted-pair) ``_id``s.
Legacy edges were keyed by ``hash("src-tgt")`` in the *call* direction,
so an edge could live under either orientation's id. After
``upsert_edge`` switched to a canonical sorted-pair id, a fresh write
lands on a different ``_id`` than a legacy reverse-direction doc,
leaving two documents for one edge (``node_degree``/``get_node_edges``
double-count). This re-keys every non-canonical doc onto its canonical
``_id`` and deletes the stale id.
**Fail-fast.** Runs in ``initialize`` inside ``get_data_init_lock``
(which serialises one deployment's worker pool — only the first worker
migrates, the rest skip via the ``_meta`` flag). On any non-benign
per-item error (e.g. 429/503) it raises, so the service does not start
until the index is fully canonical; the next startup rescans (the flag
is only set on full success). Because the service is gated on a complete
migration and every later write is canonical, there is no need for a
per-write reverse-orientation cleanup.
The canonical write uses ``op_type=create`` (insert-only): a legacy
reciprocal duplicate (both directed docs present) collapses onto the
existing forward/canonical doc (create 409, benign) and the reverse copy
is deleted. A create that fails fast happens *before* any delete, so a
source row is never dropped without its canonical counterpart existing.
Assumes no concurrent *old-version* writer adds non-canonical docs after
this completes (true for stop-the-world / single-deployment restarts). A
true rolling deploy with two code versions writing the same index could
leave a straggler reverse doc; the remedy is to clear the ``_meta`` flag
and let the next startup re-migrate.
"""
try:
if not await self.client.indices.exists(index=self._edges_index):
logger.debug(
f"[{self.workspace}] Edge index {self._edges_index} does not "
f"exist yet; skipping canonical edge-id migration"
)
return
mapping = await self.client.indices.get_mapping(index=self._edges_index)
meta = (
mapping.get(self._edges_index, {}).get("mappings", {}).get("_meta", {})
)
if meta.get(_EDGE_ID_CANONICAL_META_FLAG):
logger.info(
f"[{self.workspace}] Edge index {self._edges_index} already on "
f"canonical ids; skipping migration"
)
return
# Count upfront so operators get an X/total denominator; best-effort
# (migration still works if count is unavailable).
try:
total = (await self.client.count(index=self._edges_index)).get("count")
except OpenSearchException:
total = None
logger.info(
f"[{self.workspace}] Starting canonical edge-id migration for "
f"{self._edges_index}"
+ (f" (~{total} edges to scan)" if total is not None else "")
)
scanned = 0
migrated = 0
# Each entry is (canonical_id, old_id, source) for one non-canonical
# doc to be re-keyed. Flush roughly one bulk chunk at a time so a huge
# index does not buffer every action in memory before writing.
pending: list[tuple[str, str, dict[str, Any]]] = []
flush_at = max(self._max_upsert_records_per_batch, 1)
next_progress = _EDGE_MIGRATION_PROGRESS_INTERVAL
async def _flush_pending() -> None:
nonlocal pending
if not pending:
return
batch, pending = pending, []
# Group this batch by canonical id. A node pair can have >1
# pending non-canonical doc (e.g. several legacy reciprocal
# duplicates), and they all map to the same canonical id; carry
# each doc's old _id so it can be deleted after consolidation.
docs_by_canonical: dict[str, list[tuple[str, dict[str, Any]]]] = {}
for canonical, old_id, source in batch:
docs_by_canonical.setdefault(canonical, []).append((old_id, source))
# Phase 1 — create exactly ONE canonical doc per canonical id.
# When a canonical has >1 pending doc we pre-merge their payloads
# in memory so the single create carries the summed weight/unioned
# provenance: issuing one create per doc instead would let one win
# the insert and the rest 409, and folding those 409'd docs back in
# would re-merge the create winner (its source is now the base),
# double-counting its weight. op_type=create (insert-only): a 409
# then means the canonical already existed *independently of this
# batch* (a forward legacy doc, or a prior batch/run), so every doc
# this batch mapped to it is a loser to fold in. raise_on_error=
# False so a 409 does not abort.
create_actions = []
for canonical, reverse_docs in docs_by_canonical.items():
sources = [source for _old_id, source in reverse_docs]
source = (
{**sources[0], **_merge_edge_payloads(sources)}
if len(sources) > 1
else sources[0]
)
create_actions.append(
{
"_op_type": "create",
"_index": self._edges_index,
"_id": canonical,
"_source": source,
}
)
_success, errors = await _run_chunked_async_bulk(
self.client,
create_actions,
max_payload_bytes=self._max_upsert_payload_bytes,
max_records_per_batch=self._max_upsert_records_per_batch,
log_prefix=f"[{self.workspace}] {self.namespace} edges:",
what="canonical edge-id migration (create)",
raise_on_error=False,
)
# A create 409 means the canonical doc already exists: merge the
# pending doc(s)' relation payload into it (below) so deleting them
# loses no evidence. Any other create error (e.g. 429/503) fails
# fast BEFORE any delete, so no edge is dropped without its
# canonical counterpart in place; the flag stays unset and the
# next startup rescans.
conflicted_canonicals: set[str] = set()
real_create_errors = []
for e in errors:
info = e.get("create") if isinstance(e, dict) else None
if info is not None and info.get("status") == 409:
if info.get("_id"):
conflicted_canonicals.add(info["_id"])
continue
real_create_errors.append(e)
if real_create_errors:
raise RuntimeError(
f"Canonical edge-id migration: {len(real_create_errors)} "
f"create error(s) in {self._edges_index}; aborting startup "
f"(no source rows deleted)"
)
# The canonical pre-existed (the create did not write it), so fold
# *every* doc this batch mapped to it — none is the create winner —
# then delete their old ids (in _merge_into_canonical_edge).
for canonical in conflicted_canonicals:
await self._merge_into_canonical_edge(
canonical, docs_by_canonical[canonical]
)
# Phase 2 — every create succeeded, 409'd-then-merged, so the
# canonical now exists for all; delete the old ids. delete 404 is
# benign (another run already removed it); any other delete error
# fails fast.
delete_actions = [
{"_op_type": "delete", "_index": self._edges_index, "_id": old_id}
for _canonical, old_id, _source in batch
]
_ds, derrors = await _run_chunked_async_bulk(
self.client,
delete_actions,
max_payload_bytes=self._max_upsert_payload_bytes,
max_records_per_batch=self._max_delete_records_per_batch,
log_prefix=f"[{self.workspace}] {self.namespace} edges:",
what="canonical edge-id migration (delete)",
raise_on_error=False,
)
real_delete_errors = [
e
for e in derrors
if not (
isinstance(e, dict) and e.get("delete", {}).get("status") == 404
)
]
if real_delete_errors:
raise RuntimeError(
f"Canonical edge-id migration: {len(real_delete_errors)} "
f"delete error(s) in {self._edges_index}; aborting startup"
)
scroll_id = None
try:
response = await self.client.search(
index=self._edges_index,
body={"query": {"match_all": {}}, "sort": ["_doc"]},
scroll="5m",
size=1000,
)
while True:
scroll_id = response.get("_scroll_id")
hits = response.get("hits", {}).get("hits", [])
if not hits:
break
for hit in hits:
scanned += 1
source = hit.get("_source", {})
src = source.get("source_node_id")
tgt = source.get("target_node_id")
if not src or not tgt:
continue
canonical = _canonical_edge_id(src, tgt)
if hit["_id"] == canonical:
continue
# Queue (canonical, old_id, source); the create/delete
# split happens in _flush_pending so a failed create never
# takes its source row with it.
pending.append((canonical, hit["_id"], source))
migrated += 1
if len(pending) >= flush_at:
await _flush_pending()
if scanned >= next_progress:
logger.info(
f"[{self.workspace}] Canonical edge-id migration "
f"progress: scanned {scanned}"
+ (f"/{total}" if total is not None else "")
+ f", migrated {migrated} so far"
)
next_progress += _EDGE_MIGRATION_PROGRESS_INTERVAL
response = await self.client.scroll(
scroll_id=scroll_id, scroll="5m"
)
await _flush_pending()
finally:
if scroll_id is not None:
try:
await self.client.clear_scroll(scroll_id=scroll_id)
except OpenSearchException:
pass
if migrated:
# Make migrated docs visible to subsequent searches in one go.
try:
await self.client.indices.refresh(index=self._edges_index)
except OpenSearchException:
pass
logger.info(
f"[{self.workspace}] Canonical edge-id migration complete for "
f"{self._edges_index}: scanned {scanned}, migrated {migrated}"
)
# Mark complete (only reached on full success) so subsequent startups
# skip the full scan. Legacy reciprocal duplicates collapsed onto one
# canonical doc: the reverse doc's relation payload was merged into
# the existing canonical (see _merge_into_canonical_edge) and the
# reverse orientation deleted — no relation evidence lost.
await self.client.indices.put_mapping(
index=self._edges_index,
body={"_meta": {**meta, _EDGE_ID_CANONICAL_META_FLAG: True}},
)
except OpenSearchException as e:
# Fail fast: a transport/cluster error during migration must abort
# startup (flag stays unset) rather than serve a half-migrated index.
logger.error(
f"[{self.workspace}] Canonical edge-id migration failed for "
f"{self._edges_index}: {e}; aborting startup"
)
raise
async def _merge_into_canonical_edge(
self, canonical_id: str, reverse_docs: list[tuple[str, dict[str, Any]]]
) -> None:
"""Merge legacy reverse-orientation doc(s)' payload into an existing
canonical doc (the create-409 reciprocal-duplicate case) so deleting the
reverse loses no relation evidence (mirrors mongo_impl's dedupe merge).
``reverse_docs`` carries every pending reverse doc in this batch that maps
to ``canonical_id`` as ``(old_id, source)`` pairs (usually one, but >1
when a node pair has 3+ legacy docs); all are folded into the canonical in
a single write, then deleted by their old ids.
Uses optimistic concurrency (``if_seq_no``/``if_primary_term``) so a
concurrent live write during a rolling deploy is never clobbered: on a
version conflict we re-read — now including that write — and re-merge.
Weight summing is per-fragment and not self-idempotent (see
``_merge_edge_payloads``), so idempotency across fail-fast retries comes
from deleting each folded reverse right after the canonical write: a
re-scan no longer finds it, so it is never folded (and summed) twice. The
delete is best-effort — the batch's Phase-2 delete re-attempts the same
ids — so it deliberately does not abort the migration; the only window in
which a weight could double-count is a crash strictly between the
canonical write and this delete, a bounded and non-fatal ranking
perturbation for a one-time migration.
"""
for _attempt in range(3):
try:
current = await self.client.get(
index=self._edges_index, id=canonical_id
)
except NotFoundError:
# Canonical vanished between the create-409 and now; recreate it
# by merging the reverse sources together (nothing else to merge
# against).
reverse_sources = [source for _old_id, source in reverse_docs]
recreated = {
**reverse_sources[0],
**_merge_edge_payloads(reverse_sources),
}
await self.client.index(
index=self._edges_index, id=canonical_id, body=recreated
)
await self._delete_folded_reverse_edges(reverse_docs)
return
base = current.get("_source", {})
reverse_sources = [source for _old_id, source in reverse_docs]
merged = {**base, **_merge_edge_payloads([base, *reverse_sources])}
try:
await self.client.index(
index=self._edges_index,
id=canonical_id,
body=merged,
if_seq_no=current["_seq_no"],
if_primary_term=current["_primary_term"],
)
except ConflictError:
# A concurrent write changed the canonical doc; re-read and
# re-merge so we never overwrite that write with stale data.
continue
await self._delete_folded_reverse_edges(reverse_docs)
return
raise RuntimeError(
f"Canonical edge-id migration: could not merge into {canonical_id} "
f"after retries in {self._edges_index}; aborting startup"
)
async def _delete_folded_reverse_edges(
self, reverse_docs: list[tuple[str, dict[str, Any]]]
) -> None:
"""Delete legacy reverse docs whose payload was just folded into the
canonical, so a fail-fast re-scan never re-folds (and re-sums) them.
Best-effort: a 404 means another run already removed it, and any other
error is swallowed (logged) rather than raised — the migration's Phase-2
bulk delete re-attempts these same ids, so failing here must not abort
startup nor leave the canonical without its merged evidence.
"""
for old_id, _source in reverse_docs:
try:
await self.client.delete(index=self._edges_index, id=old_id)
except NotFoundError:
pass
except OpenSearchException as e:
logger.warning(
f"[{self.workspace}] Canonical edge-id migration: could not "
f"delete folded reverse doc {old_id} in {self._edges_index} "
f"(Phase-2 delete will retry): {e}"
)
async def finalize(self):
"""Release the OpenSearch client connection."""
if self.client is not None:
await ClientManager.release_client(self.client)
self.client = None
# --- Basic queries ---
async def has_node(self, node_id: str) -> bool:
"""Check whether a node exists in the graph."""
if not self._indices_ready:
return False
try:
return await self.client.exists(index=self._nodes_index, id=node_id)
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return False
async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
"""Check whether an edge exists between two nodes.
Startup migration is fail-fast, so after initialize() every edge is keyed
by its canonical (sorted-pair) ``_id``. Point-check that single id with
exists(), mirroring has_node() — real-time (translog-backed), independent
of the index refresh cycle, no candidate-id fan-out.
"""
if not self._indices_ready:
return False
try:
edge_id = _canonical_edge_id(source_node_id, target_node_id)
return await self.client.exists(index=self._edges_index, id=edge_id)
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return False
async def node_degree(self, node_id: str) -> int:
"""Count the number of edges connected to a node."""
if not self._indices_ready:
return 0
try:
await self._refresh_graph_indices_if_dirty(refresh_edges=True)
response = await self.client.count(
index=self._edges_index,
body={
"query": {
"bool": {
"should": [
{"term": {"source_node_id": node_id}},
{"term": {"target_node_id": node_id}},
]
}
}
},
)
return response.get("count", 0)
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return 0
async def edge_degree(self, src_id: str, tgt_id: str) -> int:
"""Sum of degrees of both endpoint nodes."""
src_degree = await self.node_degree(src_id)
tgt_degree = await self.node_degree(tgt_id)
return src_degree + tgt_degree
async def get_node(self, node_id: str) -> dict[str, str] | None:
"""Get a node document by ID, or None if not found."""
if not self._indices_ready:
return None
try:
response = await _mget_optional_doc(self.client, self._nodes_index, node_id)
if response is None:
return None
doc = response["_source"]
doc["_id"] = response["_id"]
return doc
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return None
async def get_edge(
self, source_node_id: str, target_node_id: str
) -> dict[str, str] | None:
"""Get an edge between two nodes, or None.
Edges are stored under their canonical (sorted-pair) ``_id`` once the
fail-fast startup migration completes, so read that single id directly via
mget — real-time (translog-backed), no candidate-id fan-out.
"""
if not self._indices_ready:
return None
try:
edge_id = _canonical_edge_id(source_node_id, target_node_id)
response = await _mget_optional_doc(self.client, self._edges_index, edge_id)
if response is None:
return None
result = response["_source"]
result["_id"] = response["_id"]
return result
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return None
async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None:
"""Get all (source, target) edge tuples connected to a node."""
if not self._indices_ready:
return None
try:
await self._refresh_graph_indices_if_dirty(refresh_edges=True)
query = {
"bool": {
"should": [
{"term": {"source_node_id": source_node_id}},
{"term": {"target_node_id": source_node_id}},
]
}
}
edges = []
pit = await self.client.create_pit(
index=self._edges_index, params={"keep_alive": "1m"}
)
pit_id = pit["pit_id"]
try:
search_after = None
while True:
body = {
"query": query,
"_source": ["source_node_id", "target_node_id"],
"size": 10000,
"pit": {"id": pit_id, "keep_alive": "1m"},
"sort": _pit_sort_with_composite_key(
"source_node_id", "target_node_id"
),
}
if search_after:
body["search_after"] = search_after
response = await self.client.search(body=body)
hits = response["hits"]["hits"]
if not hits:
break
for hit in hits:
edges.append(
(
hit["_source"]["source_node_id"],
hit["_source"]["target_node_id"],
)
)
search_after = hits[-1]["sort"]
if len(hits) < 10000:
break
finally:
try:
await self.client.delete_pit(body={"pit_id": [pit_id]})
except Exception:
pass
return edges
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return None
# --- Batch operations ---
async def get_nodes_batch(self, node_ids: list[str]) -> dict[str, dict]:
"""Batch-fetch multiple nodes by ID."""
if not self._indices_ready:
return {}
try:
response = await self.client.mget(
index=self._nodes_index, body={"ids": node_ids}
)
result = {}
for doc in response["docs"]:
if doc.get("found"):
data = doc["_source"]
data["_id"] = doc["_id"]
result[doc["_id"]] = data
return result
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return {}
async def node_degrees_batch(self, node_ids: list[str]) -> dict[str, int]:
"""Batch-fetch edge counts for multiple nodes using aggregations."""
if not node_ids:
return {}
if not self._indices_ready:
return {}
try:
await self._refresh_graph_indices_if_dirty(refresh_edges=True)
# Use a single query with aggregations for both source and target
body = {
"size": 0,
"query": {
"bool": {
"should": [
{"terms": {"source_node_id": node_ids}},
{"terms": {"target_node_id": node_ids}},
]
}
},
"aggs": {
"source_degrees": {
"terms": {
"field": "source_node_id",
"size": len(node_ids) * 2,
}
},
"target_degrees": {
"terms": {
"field": "target_node_id",
"size": len(node_ids) * 2,
}
},
},
}
response = await self.client.search(index=self._edges_index, body=body)
result = {}
for bucket in response["aggregations"]["source_degrees"]["buckets"]:
if bucket["key"] in node_ids:
result[bucket["key"]] = (
result.get(bucket["key"], 0) + bucket["doc_count"]
)
for bucket in response["aggregations"]["target_degrees"]["buckets"]:
if bucket["key"] in node_ids:
result[bucket["key"]] = (
result.get(bucket["key"], 0) + bucket["doc_count"]
)
return result
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return {}
async def get_nodes_edges_batch(
self, node_ids: list[str]
) -> dict[str, list[tuple[str, str]]]:
"""Batch-fetch edge tuples for multiple nodes."""
result = {nid: [] for nid in node_ids}
if not self._indices_ready:
return result
try:
await self._refresh_graph_indices_if_dirty(refresh_edges=True)
query = {
"bool": {
"should": [
{"terms": {"source_node_id": node_ids}},
{"terms": {"target_node_id": node_ids}},
]
}
}
pit = await self.client.create_pit(
index=self._edges_index, params={"keep_alive": "1m"}
)
pit_id = pit["pit_id"]
try:
search_after = None
while True:
body = {
"query": query,
"_source": ["source_node_id", "target_node_id"],
"size": 10000,
"pit": {"id": pit_id, "keep_alive": "1m"},
"sort": _pit_sort_with_composite_key(
"source_node_id", "target_node_id"
),
}
if search_after:
body["search_after"] = search_after
response = await self.client.search(body=body)
hits = response["hits"]["hits"]
if not hits:
break
for hit in hits:
src = hit["_source"]["source_node_id"]
tgt = hit["_source"]["target_node_id"]
if src in result:
result[src].append((src, tgt))
if tgt in result:
result[tgt].append((src, tgt))
search_after = hits[-1]["sort"]
if len(hits) < 10000:
break
finally:
try:
await self.client.delete_pit(body={"pit_id": [pit_id]})
except Exception:
pass
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
pass
return result
# --- Upsert operations ---
async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
"""Insert or update a node. Adds entity_id for PPL compatibility."""
try:
await self._ensure_indices_ready()
doc = {k: v for k, v in node_data.items() if k != "_id"}
doc["entity_id"] = node_id
if node_data.get("source_id", ""):
doc["source_ids"] = node_data["source_id"].split(GRAPH_FIELD_SEP)
# No per-operation refresh: node reads use ID-based mget/exists
# (translog, real-time). Search visibility after index_done_callback().
await self.client.index(index=self._nodes_index, id=node_id, body=doc)
self._nodes_dirty = True
except OpenSearchException as e:
logger.error(f"[{self.workspace}] Error upserting node {node_id}: {e}")
async def upsert_edge(
self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
) -> None:
"""Insert or update an edge keyed by a canonical (sorted-pair) ``_id``.
The canonical id collapses ``(src, tgt)`` and ``(tgt, src)`` onto one
document, so this is idempotent by construction: concurrent
reciprocal writers overwrite the same ``_id`` (last-write-wins) instead
of racing into two docs. No ``exists(reverse)`` read-then-write needed.
New writes are always canonical, and the startup migration is fail-fast
(the service does not start until every legacy doc is on its canonical
id), so there is no need to delete a reverse-orientation doc on each
write — the index is canonical before any write happens.
"""
try:
await self._ensure_indices_ready()
# Ensure source node exists (don't overwrite if it already has data)
if not await self.has_node(source_node_id):
await self.upsert_node(source_node_id, {})
doc = {k: v for k, v in edge_data.items() if k != "_id"}
doc["source_node_id"] = source_node_id
doc["target_node_id"] = target_node_id
if edge_data.get("source_id", ""):
doc["source_ids"] = edge_data["source_id"].split(GRAPH_FIELD_SEP)
edge_id = _canonical_edge_id(source_node_id, target_node_id)
await self.client.index(index=self._edges_index, id=edge_id, body=doc)
self._edges_dirty = True
except OpenSearchException as e:
logger.error(
f"[{self.workspace}] Error upserting edge {source_node_id}->{target_node_id}: {e}"
)
async def upsert_nodes_batch(self, nodes: list[tuple[str, dict[str, str]]]) -> None:
"""Batch insert/update multiple nodes using the OpenSearch bulk API.
Args:
nodes: List of (node_id, node_data) tuples.
"""
if not nodes:
return
try:
await self._ensure_indices_ready()
actions = []
for node_id, node_data in nodes:
doc = {k: v for k, v in node_data.items() if k != "_id"}
doc["entity_id"] = node_id
if node_data.get("source_id", ""):
doc["source_ids"] = node_data["source_id"].split(GRAPH_FIELD_SEP)
actions.append(
{
"_op_type": "index",
"_index": self._nodes_index,
"_id": node_id,
"_source": doc,
}
)
await _run_chunked_async_bulk(
self.client,
actions,
max_payload_bytes=self._max_upsert_payload_bytes,
max_records_per_batch=self._max_upsert_records_per_batch,
log_prefix=f"[{self.workspace}] {self.namespace} nodes:",
what="node upsert",
raise_on_error=True,
)
self._nodes_dirty = True
except OpenSearchException as e:
logger.error(f"[{self.workspace}] Error during batch node upsert: {e}")
async def has_nodes_batch(self, node_ids: list[str]) -> set[str]:
"""Check existence of multiple nodes using a single mget request.
Args:
node_ids: List of node IDs to check.
Returns:
Set of node_ids that exist in the graph.
"""
if not node_ids:
return set()
if not self._indices_ready:
return set()
try:
response = await self.client.mget(
index=self._nodes_index, body={"ids": node_ids}
)
return {doc["_id"] for doc in response.get("docs", []) if doc.get("found")}
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return set()
async def upsert_edges_batch(
self, edges: list[tuple[str, str, dict[str, str]]]
) -> None:
"""Batch insert/update multiple edges using the OpenSearch bulk API.
Each edge is keyed by its canonical (sorted-pair) ``_id`` (see
``_canonical_edge_id``), so reciprocal directions collapse onto one
document with no reverse-direction look-up. Edges that map to the same
canonical id within this batch are deduplicated last-write-wins.
Args:
edges: List of (source_node_id, target_node_id, edge_data) tuples.
"""
if not edges:
return
try:
await self._ensure_indices_ready()
# Ensure all source nodes exist (mirrors upsert_edge behaviour)
source_ids = list({src for src, _tgt, _data in edges})
existing_sources = await self.has_nodes_batch(source_ids)
missing_sources = [
(nid, {}) for nid in source_ids if nid not in existing_sources
]
if missing_sources:
await self.upsert_nodes_batch(missing_sources)
# Key every edge by its canonical id and dedupe within the batch
# (last-write-wins) so a single bulk request carries one action per
# logical edge regardless of direction.
actions_by_id: dict[str, dict[str, Any]] = {}
for src, tgt, edge_data in edges:
doc = {k: v for k, v in edge_data.items() if k != "_id"}
doc["source_node_id"] = src
doc["target_node_id"] = tgt
if edge_data.get("source_id", ""):
doc["source_ids"] = edge_data["source_id"].split(GRAPH_FIELD_SEP)
edge_id = _canonical_edge_id(src, tgt)
actions_by_id[edge_id] = {
"_op_type": "index",
"_index": self._edges_index,
"_id": edge_id,
"_source": doc,
}
actions = list(actions_by_id.values())
await _run_chunked_async_bulk(
self.client,
actions,
max_payload_bytes=self._max_upsert_payload_bytes,
max_records_per_batch=self._max_upsert_records_per_batch,
log_prefix=f"[{self.workspace}] {self.namespace} edges:",
what="edge upsert",
raise_on_error=True,
)
self._edges_dirty = True
except OpenSearchException as e:
logger.error(f"[{self.workspace}] Error during batch edge upsert: {e}")
# --- Delete operations ---
async def delete_node(self, node_id: str) -> None:
"""Delete a node and all its connected edges.
Marks node and edge search views dirty so refresh happens lazily on the
next search/count-based graph read. Uses conflicts="proceed" to
tolerate already-deleted matches.
"""
try:
# Refresh edge search view so delete_by_query sees all un-flushed writes.
await self._refresh_graph_indices_if_dirty(refresh_edges=True)
# Delete all edges referencing this node
body = {
"query": {
"bool": {
"should": [
{"term": {"source_node_id": node_id}},
{"term": {"target_node_id": node_id}},
]
}
}
}
await self.client.delete_by_query(
index=self._edges_index,
body=body,
params={"conflicts": "proceed"},
)
# Delete the node
try:
await self.client.delete(index=self._nodes_index, id=node_id)
except NotFoundError:
pass
self._nodes_dirty = True
self._edges_dirty = True
except OpenSearchException as e:
logger.error(f"[{self.workspace}] Error deleting node {node_id}: {e}")
async def remove_nodes(self, nodes: list[str]) -> None:
"""Batch-delete multiple nodes and their connected edges.
Marks node and edge search views dirty so refresh happens lazily on the
next search/count-based graph read. Uses conflicts="proceed" to
tolerate already-deleted matches.
"""
if not nodes:
return
logger.info(f"[{self.workspace}] Deleting {len(nodes)} nodes")
try:
# Refresh edge search view so delete_by_query sees all un-flushed writes.
await self._refresh_graph_indices_if_dirty(refresh_edges=True)
# Delete edges
body = {
"query": {
"bool": {
"should": [
{"terms": {"source_node_id": nodes}},
{"terms": {"target_node_id": nodes}},
]
}
}
}
await self.client.delete_by_query(
index=self._edges_index,
body=body,
params={"conflicts": "proceed"},
)
# Delete nodes
actions = [
{"_op_type": "delete", "_index": self._nodes_index, "_id": nid}
for nid in nodes
]
await _run_chunked_async_bulk(
self.client,
actions,
max_payload_bytes=self._max_upsert_payload_bytes,
max_records_per_batch=self._max_delete_records_per_batch,
log_prefix=f"[{self.workspace}] {self.namespace} nodes:",
what="node delete",
raise_on_error=False,
)
self._nodes_dirty = True
self._edges_dirty = True
except OpenSearchException as e:
logger.error(f"[{self.workspace}] Error removing nodes: {e}")
async def remove_edges(self, edges: list[tuple[str, str]]) -> None:
"""Batch-delete multiple edges by canonical ID (real-time).
Startup migration is fail-fast, so every edge is keyed by its canonical
(sorted-pair) ``_id``. Delete that single id per edge — dedup via a set so
reciprocal inputs ``(A,B)`` and ``(B,A)`` collapse to one delete op. The
raw bulk API does not raise on a 404 delete.
Marks edge search views dirty so refresh happens lazily on the next
search/count-based graph read.
"""
if not edges:
return
logger.info(f"[{self.workspace}] Deleting {len(edges)} edges")
try:
edge_ids = {_canonical_edge_id(src, tgt) for src, tgt in edges}
operations = [
{
"delete": {
"_index": self._edges_index,
"_id": edge_id,
}
}
for edge_id in edge_ids
]
# The raw bulk API does not auto-chunk (unlike helpers.async_bulk),
# so split the operation list by the delete record-count cap to keep
# each request bounded (mirrors mongo_impl's chunked delete_many).
chunk = (
self._max_delete_records_per_batch
if self._max_delete_records_per_batch > 0
else len(operations)
)
if len(operations) > chunk:
logger.info(
f"[{self.workspace}] {self.namespace} edges: edge delete "
f"{len(operations)} ops split into bulk chunks (chunk={chunk})"
)
for i in range(0, len(operations), chunk):
await self.client.bulk(body=operations[i : i + chunk])
self._edges_dirty = True
except OpenSearchException as e:
logger.error(f"[{self.workspace}] Error removing edges: {e}")
# --- Query operations ---
async def get_all_labels(self) -> list[str]:
"""Get all node IDs (entity names) sorted alphabetically."""
if not self._indices_ready:
return []
try:
await self._refresh_graph_indices_if_dirty(refresh_nodes=True)
labels = []
pit = await self.client.create_pit(
index=self._nodes_index, params={"keep_alive": "1m"}
)
pit_id = pit["pit_id"]
try:
search_after = None
while True:
body = {
"query": {"match_all": {}},
"_source": False,
"size": 10000,
"pit": {"id": pit_id, "keep_alive": "1m"},
"sort": _pit_sort_with_field("entity_id"),
}
if search_after:
body["search_after"] = search_after
response = await self.client.search(body=body)
hits = response["hits"]["hits"]
if not hits:
break
for hit in hits:
labels.append(hit["_id"])
search_after = hits[-1]["sort"]
if len(hits) < 10000:
break
finally:
try:
await self.client.delete_pit(body={"pit_id": [pit_id]})
except Exception:
pass
labels.sort()
return labels
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return []
async def _collect_node_ids(
self, limit: int, exclude_ids: set[str] | None = None
) -> list[str]:
"""Collect up to `limit` node IDs, optionally skipping known IDs."""
if limit <= 0:
return []
excluded = exclude_ids or set()
if not excluded and limit <= 10000:
body = {
"query": {"match_all": {}},
"_source": False,
"size": limit,
}
resp = await self.client.search(index=self._nodes_index, body=body)
return [hit["_id"] for hit in resp["hits"]["hits"]]
node_ids: list[str] = []
pit = await self.client.create_pit(
index=self._nodes_index, params={"keep_alive": "1m"}
)
pit_id = pit["pit_id"]
try:
search_after = None
while len(node_ids) < limit:
body = {
"query": {"match_all": {}},
"_source": False,
"size": 10000,
"pit": {"id": pit_id, "keep_alive": "1m"},
"sort": _pit_sort_with_field("entity_id"),
}
if search_after:
body["search_after"] = search_after
resp = await self.client.search(body=body)
hits = resp["hits"]["hits"]
if not hits:
break
for hit in hits:
node_id = hit["_id"]
if node_id in excluded:
continue
node_ids.append(node_id)
if len(node_ids) >= limit:
break
search_after = hits[-1].get("sort")
if len(hits) < 10000:
break
finally:
try:
await self.client.delete_pit(body={"pit_id": [pit_id]})
except Exception:
pass
return node_ids
@staticmethod
def _edge_rank_key(edge: dict[str, Any]) -> tuple[int, float]:
"""Rank traversal edges by shallower depth first, then higher weight."""
depth = edge.get("_depth", edge.get("depth", 0))
try:
depth_value = int(depth)
except (TypeError, ValueError):
depth_value = 0
weight = edge.get("weight", 0)
try:
weight_value = float(weight)
except (TypeError, ValueError):
weight_value = 0.0
return (depth_value, -weight_value)
async def _append_edges_between_nodes(
self, node_ids: list[str], result: KnowledgeGraph
) -> None:
"""Append all edges whose source and target are both in `node_ids`."""
if not node_ids:
return
edge_query = {
"bool": {
"must": [
{"terms": {"source_node_id": node_ids}},
{"terms": {"target_node_id": node_ids}},
]
}
}
seen_edges = set()
pit = await self.client.create_pit(
index=self._edges_index, params={"keep_alive": "1m"}
)
pit_id = pit["pit_id"]
try:
search_after = None
while True:
edge_body = {
"query": edge_query,
"size": 10000,
"pit": {"id": pit_id, "keep_alive": "1m"},
"sort": _pit_sort_with_composite_key(
"source_node_id", "target_node_id"
),
}
if search_after:
edge_body["search_after"] = search_after
edge_resp = await self.client.search(body=edge_body)
hits = edge_resp["hits"]["hits"]
if not hits:
break
for hit in hits:
e = hit["_source"]
eid = f"{e['source_node_id']}-{e['target_node_id']}"
if eid not in seen_edges:
seen_edges.add(eid)
result.edges.append(self._construct_graph_edge(eid, e))
search_after = hits[-1].get("sort")
if len(hits) < 10000:
break
finally:
try:
await self.client.delete_pit(body={"pit_id": [pit_id]})
except Exception:
pass
def _construct_graph_node(self, node_id, node_data: dict) -> KnowledgeGraphNode:
return KnowledgeGraphNode(
id=node_id,
labels=[node_id],
properties={
k: v
for k, v in node_data.items()
if k
not in (
"_id",
"source_ids",
"connected_edges",
"edge_count",
)
},
)
def _construct_graph_edge(self, edge_id: str, edge: dict) -> KnowledgeGraphEdge:
return KnowledgeGraphEdge(
id=edge_id,
type=edge.get("relationship", ""),
source=edge["source_node_id"],
target=edge["target_node_id"],
properties={
k: v
for k, v in edge.items()
if k
not in (
"_id",
"source_node_id",
"target_node_id",
"relationship",
"source_ids",
)
},
)
async def get_knowledge_graph(
self,
node_label: str,
max_depth: int = 3,
max_nodes: int = None,
) -> KnowledgeGraph:
"""Retrieve a subgraph via PPL graphlookup (if available) or client-side BFS."""
if not self._indices_ready:
return KnowledgeGraph()
if max_nodes is None:
max_nodes = self.global_config.get("max_graph_nodes", 1000)
else:
max_nodes = min(max_nodes, self.global_config.get("max_graph_nodes", 1000))
result = KnowledgeGraph()
start = time.perf_counter()
try:
await self._refresh_graph_indices_if_dirty(
refresh_nodes=True, refresh_edges=True
)
if node_label == "*":
result = await self._get_knowledge_graph_all(max_nodes)
elif self._ppl_graphlookup_available:
result = await self._bfs_subgraph_ppl(node_label, max_depth, max_nodes)
else:
result = await self._bfs_subgraph(node_label, max_depth, max_nodes)
duration = time.perf_counter() - start
logger.info(
f"[{self.workspace}] Subgraph query in {duration:.4f}s | "
f"Nodes: {len(result.nodes)} | Edges: {len(result.edges)} | Truncated: {result.is_truncated}"
)
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return KnowledgeGraph()
logger.error(f"[{self.workspace}] Graph query failed: {e}")
return result
async def _get_knowledge_graph_all(self, max_nodes: int) -> KnowledgeGraph:
"""Get all nodes (up to max_nodes, ranked by degree) and their interconnecting edges."""
result = KnowledgeGraph()
if not self._indices_ready:
return result
try:
total = (await self.client.count(index=self._nodes_index))["count"]
result.is_truncated = total > max_nodes
if result.is_truncated:
# Get top nodes by degree
body = {
"size": 0,
"aggs": {
"src": {
"terms": {
"field": "source_node_id",
"size": max_nodes,
}
},
"tgt": {
"terms": {
"field": "target_node_id",
"size": max_nodes,
}
},
},
}
resp = await self.client.search(index=self._edges_index, body=body)
degree_map = {}
for bucket in resp["aggregations"]["src"]["buckets"]:
degree_map[bucket["key"]] = (
degree_map.get(bucket["key"], 0) + bucket["doc_count"]
)
for bucket in resp["aggregations"]["tgt"]["buckets"]:
degree_map[bucket["key"]] = (
degree_map.get(bucket["key"], 0) + bucket["doc_count"]
)
top_ids = sorted(degree_map, key=degree_map.get, reverse=True)[
:max_nodes
]
if len(top_ids) < max_nodes:
top_ids.extend(
await self._collect_node_ids(
max_nodes - len(top_ids), exclude_ids=set(top_ids)
)
)
else:
top_ids = await self._collect_node_ids(max_nodes)
# Fetch node data
if top_ids:
node_resp = await self.client.mget(
index=self._nodes_index, body={"ids": top_ids}
)
found_node_ids = []
for doc in node_resp["docs"]:
if doc.get("found"):
found_node_ids.append(doc["_id"])
result.nodes.append(
self._construct_graph_node(doc["_id"], doc["_source"])
)
await self._append_edges_between_nodes(found_node_ids, result)
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return result
logger.error(f"[{self.workspace}] Error in get_knowledge_graph_all: {e}")
return result
async def _bfs_subgraph_ppl(
self, start_label: str, max_depth: int, max_nodes: int
) -> KnowledgeGraph:
"""Server-side BFS using PPL graphlookup command.
Queries the nodes index for the start node, then uses graphLookup to traverse
the edges index with bidirectional BFS. Uses `flatten` to unnest results and
`depthField` for depth-based sorting. Falls back to client-side BFS on failure.
"""
result = KnowledgeGraph()
# Verify start node exists
start_node = await self.get_node(start_label)
if not start_node:
return result
result.nodes.append(self._construct_graph_node(start_label, start_node))
if max_depth == 0:
return result
# PPL maxDepth=0 means 1 hop (direct match), so max_depth-1
ppl_depth = max(0, max_depth - 1)
escaped = self._escape_ppl(start_label)
ppl_query = (
f"source = {self._nodes_index}"
f" | where entity_id = '{escaped}'"
f" | graphLookup {self._edges_index}"
f" start=entity_id"
f" edge=target_node_id<->source_node_id"
f" maxDepth={ppl_depth}"
f" depthField=_depth"
f" usePIT=true"
f" as connected_edges"
)
try:
resp = await self.client.transport.perform_request(
"POST",
"/_plugins/_ppl",
body={"query": ppl_query},
)
except Exception as e:
logger.warning(
f"[{self.workspace}] PPL graphlookup failed, falling back to client BFS: {e}"
)
return await self._bfs_subgraph(start_label, max_depth, max_nodes)
# Parse PPL response — schema-driven to avoid fragile positional access
try:
datarows = resp.get("datarows", [])
schema = [col["name"] for col in resp.get("schema", [])]
ce_idx = (
schema.index("connected_edges") if "connected_edges" in schema else -1
)
# Collect all edge rows from connected_edges arrays
all_edge_rows = []
for row in datarows:
edges_arr = row[ce_idx] if ce_idx >= 0 else []
if isinstance(edges_arr, list):
all_edge_rows.extend(edges_arr)
if not all_edge_rows:
return result
if isinstance(all_edge_rows[0], dict):
sorted_edge_rows = sorted(all_edge_rows, key=self._edge_rank_key)
else:
# Positional array — column positions are unknown, fall back to client BFS
logger.warning(
f"[{self.workspace}] PPL returned positional arrays, falling back to client BFS"
)
return await self._bfs_subgraph(start_label, max_depth, max_nodes)
except (KeyError, IndexError, TypeError, ValueError) as e:
logger.warning(
f"[{self.workspace}] Error parsing PPL response, falling back: {e}"
)
return await self._bfs_subgraph(start_label, max_depth, max_nodes)
ordered_node_ids = [start_label]
discovered_nodes = {start_label}
for edge_row in sorted_edge_rows:
for node_id in (
edge_row.get("source_node_id"),
edge_row.get("target_node_id"),
):
if not node_id or node_id in discovered_nodes:
continue
discovered_nodes.add(node_id)
if len(ordered_node_ids) < max_nodes:
ordered_node_ids.append(node_id)
result.is_truncated = len(discovered_nodes) > max_nodes
# Batch fetch node data (start node already added)
new_node_ids = [nid for nid in ordered_node_ids if nid != start_label]
if new_node_ids:
node_resp = await self.client.mget(
index=self._nodes_index, body={"ids": new_node_ids}
)
for doc in node_resp["docs"]:
if doc.get("found"):
result.nodes.append(
self._construct_graph_node(doc["_id"], doc["_source"])
)
await self._append_edges_between_nodes(ordered_node_ids, result)
return result
@staticmethod
def _escape_ppl(value: str) -> str:
"""Escape a string for safe inclusion in a PPL single-quoted literal.
Escapes backslashes, single quotes, and control characters that could
interfere with PPL query parsing.
"""
value = value.replace("\\", "\\\\").replace("'", "\\'")
# Strip control characters that could break the PPL string literal
value = value.replace("\n", " ").replace("\r", " ").replace("\t", " ")
return value
@staticmethod
def _escape_wildcard(value: str) -> str:
"""Escape OpenSearch wildcard special characters in user input.
Escapes \\, *, and ? so they are treated as literal characters
rather than wildcard operators, preventing DoS via expensive patterns.
"""
# Escape backslash first, then wildcard metacharacters
return value.replace("\\", "\\\\").replace("*", "\\*").replace("?", "\\?")
async def _bfs_subgraph(
self, start_label: str, max_depth: int, max_nodes: int
) -> KnowledgeGraph:
"""BFS traversal from a starting node, batching neighbor lookups per level."""
result = KnowledgeGraph()
seen_nodes = set()
# Verify start node exists
start_node = await self.get_node(start_label)
if not start_node:
return result
seen_nodes.add(start_label)
result.nodes.append(self._construct_graph_node(start_label, start_node))
current_level = [start_label]
for _ in range(max_depth):
if not current_level or len(seen_nodes) >= max_nodes:
break
# Batch fetch all edges for current level
body = {
"query": {
"bool": {
"should": [
{"terms": {"source_node_id": current_level}},
{"terms": {"target_node_id": current_level}},
]
}
},
"_source": ["source_node_id", "target_node_id"],
"size": 10000,
}
try:
resp = await self.client.search(index=self._edges_index, body=body)
except OpenSearchException:
break
next_level = set()
for hit in resp["hits"]["hits"]:
src = hit["_source"]["source_node_id"]
tgt = hit["_source"]["target_node_id"]
if src not in seen_nodes:
next_level.add(src)
if tgt not in seen_nodes:
next_level.add(tgt)
# Limit to max_nodes
new_ids = []
for nid in next_level:
if len(seen_nodes) + len(new_ids) >= max_nodes:
break
new_ids.append(nid)
if new_ids:
# Batch fetch node data
node_resp = await self.client.mget(
index=self._nodes_index, body={"ids": new_ids}
)
for doc in node_resp["docs"]:
if doc.get("found"):
seen_nodes.add(doc["_id"])
result.nodes.append(
self._construct_graph_node(doc["_id"], doc["_source"])
)
current_level = new_ids
# Fetch all edges between seen nodes using PIT scrolling
all_ids = list(seen_nodes)
if all_ids:
try:
await self._append_edges_between_nodes(all_ids, result)
except OpenSearchException:
pass
result.is_truncated = len(seen_nodes) >= max_nodes
return result
async def get_all_nodes(self) -> list[dict]:
"""Get all nodes with their properties."""
if not self._indices_ready:
return []
try:
await self._refresh_graph_indices_if_dirty(refresh_nodes=True)
nodes = []
pit = await self.client.create_pit(
index=self._nodes_index, params={"keep_alive": "1m"}
)
pit_id = pit["pit_id"]
try:
search_after = None
while True:
body = {
"query": {"match_all": {}},
"size": 10000,
"pit": {"id": pit_id, "keep_alive": "1m"},
"sort": _pit_sort_with_field("entity_id"),
}
if search_after:
body["search_after"] = search_after
response = await self.client.search(body=body)
hits = response["hits"]["hits"]
if not hits:
break
for hit in hits:
node = hit["_source"]
node["id"] = hit["_id"]
nodes.append(node)
search_after = hits[-1]["sort"]
if len(hits) < 10000:
break
finally:
try:
await self.client.delete_pit(body={"pit_id": [pit_id]})
except Exception:
pass
return nodes
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return []
async def get_all_edges(self) -> list[dict]:
"""Get all edges with source/target fields added."""
if not self._indices_ready:
return []
try:
await self._refresh_graph_indices_if_dirty(refresh_edges=True)
edges = []
pit = await self.client.create_pit(
index=self._edges_index, params={"keep_alive": "1m"}
)
pit_id = pit["pit_id"]
try:
search_after = None
while True:
body = {
"query": {"match_all": {}},
"size": 10000,
"pit": {"id": pit_id, "keep_alive": "1m"},
"sort": _pit_sort_with_composite_key(
"source_node_id", "target_node_id"
),
}
if search_after:
body["search_after"] = search_after
response = await self.client.search(body=body)
hits = response["hits"]["hits"]
if not hits:
break
for hit in hits:
edge = hit["_source"]
edge["source"] = edge.get("source_node_id")
edge["target"] = edge.get("target_node_id")
edges.append(edge)
search_after = hits[-1]["sort"]
if len(hits) < 10000:
break
finally:
try:
await self.client.delete_pit(body={"pit_id": [pit_id]})
except Exception:
pass
return edges
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return []
async def get_popular_labels(self, limit: int = 300) -> list[str]:
"""Get node labels ranked by edge degree (most connected first)."""
if not self._indices_ready:
return []
try:
await self._refresh_graph_indices_if_dirty(refresh_edges=True)
body = {
"size": 0,
"aggs": {
"src": {"terms": {"field": "source_node_id", "size": limit * 2}},
"tgt": {"terms": {"field": "target_node_id", "size": limit * 2}},
},
}
response = await self.client.search(index=self._edges_index, body=body)
degree_map = {}
for bucket in response["aggregations"]["src"]["buckets"]:
degree_map[bucket["key"]] = (
degree_map.get(bucket["key"], 0) + bucket["doc_count"]
)
for bucket in response["aggregations"]["tgt"]["buckets"]:
degree_map[bucket["key"]] = (
degree_map.get(bucket["key"], 0) + bucket["doc_count"]
)
sorted_labels = sorted(degree_map, key=degree_map.get, reverse=True)[:limit]
return sorted_labels
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return []
async def search_labels(self, query: str, limit: int = 50) -> list[str]:
"""Search node labels with wildcard and prefix matching."""
query = query.strip()
if not query:
return []
if not self._indices_ready:
return []
try:
await self._refresh_graph_indices_if_dirty(refresh_nodes=True)
body = {
"query": {
"bool": {
"should": [
{"term": {"entity_id": {"value": query, "boost": 10}}},
{
"prefix": {
"entity_id": {"value": query.lower(), "boost": 5}
}
},
{
"wildcard": {
"entity_id": {
"value": f"*{self._escape_wildcard(query.lower())}*",
"case_insensitive": True,
"boost": 2,
}
}
},
]
}
},
"_source": False,
"size": limit,
}
response = await self.client.search(index=self._nodes_index, body=body)
return [hit["_id"] for hit in response["hits"]["hits"]]
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return []
async def index_done_callback(self) -> None:
"""Refresh both node and edge indices."""
if not self._indices_ready:
return
try:
await self._refresh_graph_indices_if_dirty(
refresh_nodes=True, refresh_edges=True
)
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_indices_missing()
return
raise
async def drop(self) -> dict[str, str]:
"""Delete both node and edge indices."""
errors = []
for idx in (self._nodes_index, self._edges_index):
try:
await self.client.indices.delete(index=idx)
logger.info(f"[{self.workspace}] Dropped graph index: {idx}")
except NotFoundError:
logger.info(
f"[{self.workspace}] Graph index already missing during drop: {idx}"
)
except OpenSearchException as e:
errors.append(f"{idx}: {e}")
logger.error(
f"[{self.workspace}] Error dropping graph index {idx}: {e}"
)
except Exception as e:
errors.append(f"{idx}: {e}")
logger.error(
f"[{self.workspace}] Unexpected error dropping graph index {idx}: {e}"
)
self._mark_indices_missing()
if errors:
return {
"status": "error",
"message": "Failed to drop graph indices: " + "; ".join(errors),
}
try:
logger.info(f"[{self.workspace}] Dropped graph indices")
return {"status": "success", "message": "Graph indices dropped"}
except Exception as e:
logger.error(f"[{self.workspace}] Error finalizing graph drop: {e}")
return {"status": "error", "message": str(e)}
@final
@dataclass
class OpenSearchVectorDBStorage(BaseVectorStorage):
"""Vector storage using OpenSearch k-NN plugin with corrected cosine score handling."""
client: AsyncOpenSearch = field(default=None)
_index_name: str = field(default="", init=False)
_index_ready: bool = field(default=False, init=False)
def __init__(
self, namespace, global_config, embedding_func, workspace=None, meta_fields=None
):
super().__init__(
namespace=namespace,
workspace=workspace or "",
global_config=global_config,
embedding_func=embedding_func,
meta_fields=meta_fields or set(),
)
self.__post_init__()
def __post_init__(self):
validate_workspace(self.workspace)
self._validate_embedding_func()
self.workspace, self.final_namespace, self._index_name = _build_index_name(
self.workspace, self.namespace
)
kwargs = self.global_config.get("vector_db_storage_cls_kwargs", {})
cosine_threshold = kwargs.get("cosine_better_than_threshold")
if cosine_threshold is None:
raise ValueError(
"cosine_better_than_threshold must be specified in vector_db_storage_cls_kwargs"
)
self.cosine_better_than_threshold = cosine_threshold
self._max_batch_size = self.global_config["embedding_batch_num"]
# Pending writes are flushed via _flush_pending_vector_ops() during
# index_done_callback() / finalize(). This batches many small upsert()
# invocations into a single async_bulk roundtrip. See issue #2785.
self._pending_vector_docs: dict[str, _PendingVectorDoc] = {}
self._pending_vector_deletes: set[str] = set()
# Namespace-keyed lock (multi-process safe) is initialised in
# initialize(). All buffer reads / writes and any destructive server
# mutation (delete_by_query, drop, finalize) are serialised through
# this lock to keep in-process readers race-free during a flush and
# to order cross-worker flushes against the same OpenSearch index.
self._flush_lock = None
(
self._max_upsert_payload_bytes,
self._max_upsert_records_per_batch,
self._max_delete_records_per_batch,
) = _resolve_bulk_batch_limits()
async def initialize(self):
"""Initialize client and create k-NN vector index."""
async with get_data_init_lock():
if self.client is None:
self.client = await ClientManager.get_client()
await self._create_knn_index_if_not_exists()
self._index_ready = True
logger.debug(
f"[{self.workspace}] OpenSearch Vector storage initialized: {self._index_name}"
)
if self._flush_lock is None:
self._flush_lock = get_namespace_lock(
self.namespace, workspace=self.workspace
)
async def _ensure_index_ready(self):
"""Recreate the vector index before the next write if it is missing."""
if self._index_ready:
return
async with get_data_init_lock():
if self.client is None:
self.client = await ClientManager.get_client()
if not self._index_ready:
await self._create_knn_index_if_not_exists()
self._index_ready = True
def _mark_index_missing(self):
"""Mark the vector index as unavailable for subsequent read short-circuiting."""
self._index_ready = False
async def _create_knn_index_if_not_exists(self):
try:
if await self.client.indices.exists(index=self._index_name):
# Validate existing index dimension
try:
mapping = await self.client.indices.get_mapping(
index=self._index_name
)
existing_dim = (
mapping[self._index_name]["mappings"]["properties"]
.get("vector", {})
.get("dimension")
)
expected_dim = self.embedding_func.embedding_dim
if existing_dim is not None and existing_dim != expected_dim:
raise ValueError(
f"Vector dimension mismatch! Index '{self._index_name}' has "
f"dimension {existing_dim}, but current embedding model expects "
f"dimension {expected_dim}. Please drop the existing index or "
f"use an embedding model with matching dimensions."
)
except (KeyError, TypeError):
logger.warning(
f"[{self.workspace}] Could not read vector mapping for index "
f"'{self._index_name}'; skipping dimension validation"
)
return
ef_construction = int(
_get_opensearch_env("OPENSEARCH_KNN_EF_CONSTRUCTION", "200")
)
m = int(_get_opensearch_env("OPENSEARCH_KNN_M", "16"))
ef_search = int(_get_opensearch_env("OPENSEARCH_KNN_EF_SEARCH", "100"))
body = {
"settings": {
"index": {
"knn": True,
"knn.algo_param.ef_search": ef_search,
"number_of_shards": _get_index_number_of_shards(),
"number_of_replicas": _get_index_number_of_replicas(),
}
},
"mappings": {
"properties": {
"vector": {
"type": "knn_vector",
"dimension": self.embedding_func.embedding_dim,
"method": {
"name": "hnsw",
"space_type": "cosinesimil",
"engine": "lucene",
"parameters": {
"ef_construction": ef_construction,
"m": m,
},
},
},
"content": {"type": "text"},
"entity_name": {"type": "keyword"},
"src_id": {"type": "keyword"},
"tgt_id": {"type": "keyword"},
"file_path": {"type": "keyword"},
"created_at": {"type": "long"},
},
"dynamic": True,
},
}
await self.client.indices.create(index=self._index_name, body=body)
logger.info(
f"[{self.workspace}] Created k-NN index: {self._index_name} "
f"(dim={self.embedding_func.embedding_dim})"
)
except RequestError as e:
if "resource_already_exists_exception" not in str(e):
logger.error(f"[{self.workspace}] Error creating k-NN index: {e}")
raise
except OpenSearchException as e:
logger.error(f"[{self.workspace}] Error creating k-NN index: {e}")
raise
async def finalize(self):
"""Flush pending writes and release the OpenSearch client connection.
Regular flush failures (any ``Exception``) are captured so they
can be re-surfaced as a ``RuntimeError`` that names the unflushed
buffer counts -- otherwise ``LightRAG.finalize_storages()`` would
log the storage as successfully finalized while writes silently
failed to reach OpenSearch.
``BaseException`` subclasses other than ``Exception`` (notably
``asyncio.CancelledError`` / ``KeyboardInterrupt`` / ``SystemExit``)
are NOT caught: they propagate through the ``finally`` block so
shutdown cancellation is honoured and not silently swallowed.
The client is released in ``finally`` so it does not leak whether
the flush succeeded, failed, or was cancelled.
"""
flush_error: Exception | None = None
try:
try:
await self._flush_pending_vector_ops()
except Exception as e:
flush_error = e
finally:
if self.client is not None:
await ClientManager.release_client(self.client)
self.client = None
pending_docs = len(self._pending_vector_docs)
pending_deletes = len(self._pending_vector_deletes)
if flush_error is not None:
raise RuntimeError(
f"[{self.workspace}] OpenSearchVectorDBStorage.finalize() "
f"flush raised; {pending_docs} pending upserts and "
f"{pending_deletes} pending deletes were left buffered "
f"(client released, data lost)"
) from flush_error
if pending_docs or pending_deletes:
raise RuntimeError(
f"[{self.workspace}] OpenSearchVectorDBStorage.finalize() "
f"left {pending_docs} pending upserts and {pending_deletes} "
f"pending deletes buffered after final flush attempt "
f"(transient bulk failure); these writes have been lost"
)
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
"""Buffer vector docs for embedding and batched flush.
Docs are buffered in ``self._pending_vector_docs`` and flushed in a
single ``async_bulk`` call during ``index_done_callback()`` /
``finalize()``. This is a behavioral change relative to per-call
``async_bulk``: writes are not durable in OpenSearch until the next
flush, which matches the contract used by other LightRAG storage
backends ("changes will be persisted during the next
index_done_callback").
Embedding is deferred to the flush path so repeated upserts of the
same id and many small upsert calls can be embedded once in a single
batch. Flush holds the namespace lock while embedding and bulk
indexing so cross-worker destructive mutations cannot interleave with
partially-flushed vector writes.
"""
if not data:
return
await self._ensure_index_ready()
logger.debug(
f"[{self.workspace}] Buffering {len(data)} vectors for {self.namespace}"
)
current_time = int(time.time())
pending_docs: list[tuple[str, _PendingVectorDoc]] = []
for i, (k, v) in enumerate(data.items(), start=1):
content = v["content"]
source = {
"created_at": current_time,
**{k1: v1 for k1, v1 in v.items() if k1 in self.meta_fields},
}
pending_docs.append(
(
k,
_PendingVectorDoc(
source=source,
content=content,
),
)
)
await _cooperative_yield(i)
# Buffer: an upsert overrides a pending delete on the same id.
async with self._flush_lock:
for doc_id, pending_doc in pending_docs:
self._pending_vector_deletes.discard(doc_id)
self._pending_vector_docs[doc_id] = pending_doc
async def _flush_pending_vector_ops(self) -> None:
"""Flush buffered vector upserts and deletes via a single async_bulk call.
Concurrency contract: the entire flush, including deferred embedding,
runs under ``_flush_lock`` (a ``get_namespace_lock`` instance), and so
do all buffer reads / writes and destructive server mutations on this
storage. That keeps the operation sequential within the process and
orders concurrent cross-worker flushes against the same OpenSearch
index.
Embedding deliberately runs *inside* this lock (not in ``upsert`` or
lock-free): it makes deferred embedding and bulk indexing atomic
against concurrent upserts and destructive mutations (``drop`` /
``delete_entity_relation``). This is what lets
``index_done_callback`` / ``finalize`` promise that every buffered
vector is embedded and persisted on return. Moving embedding out of
the lock to avoid blocking reads would let a destructive op
interleave between embed and bulk and resurrect or drop vectors out
of order -- do not do it.
Failure handling:
* If ``_ensure_index_ready`` raises, the buffers are left intact
and the next flush will retry.
* If embedding raises, the buffers are left intact and the next
flush will retry. Model providers already retry internally, so
this is treated like a persistence failure.
* If ``async_bulk`` itself raises (network / parse error), the
buffers are left intact and the next flush will retry. Index
ops are idempotent on ``_id`` and a re-issued delete on a
missing doc is filtered out as 404 by ``_extract_bulk_failed_ids``.
* Per-doc retryable failures (408 / 429 / 5xx) stay in the
buffer for the next flush.
* Per-doc non-retryable failures (most 4xx, mapping errors) are
cleared from the buffer and logged with a sample of
(op, id, status, error) so operators can diagnose them.
"""
async with self._flush_lock:
if not self._pending_vector_docs and not self._pending_vector_deletes:
return
if self.client is None:
return
# If the index disappeared between writes (e.g. read path
# marked it missing), recreate it now. Failure leaves the
# buffers untouched and bubbles up to the caller.
await self._ensure_index_ready()
pending_docs = self._pending_vector_docs
pending_deletes = self._pending_vector_deletes
docs_to_embed = [
(doc_id, pending_doc)
for doc_id, pending_doc in pending_docs.items()
if pending_doc.vector is None
]
if docs_to_embed:
contents = [pending_doc.content for _, pending_doc in docs_to_embed]
batches = [
contents[i : i + self._max_batch_size]
for i in range(0, len(contents), self._max_batch_size)
]
# TEMP diagnostic (remove later): confirm deferred batching is
# actually coalescing per-id upserts. defer working -> docs >>
# batches; eager/per-id -> docs == batches == 1 every flush.
logger.info(
f"[{self.workspace}] {self.namespace} flush: embedding "
f"{len(docs_to_embed)} vectors in {len(batches)} batch(es) "
f"(batch_num={self._max_batch_size})"
)
try:
embeddings_list = await asyncio.gather(
*[
self.embedding_func(batch, context="document")
for batch in batches
]
)
except Exception as e:
logger.error(
f"[{self.workspace}] Error embedding pending vector ops "
f"(upserts={len(docs_to_embed)}): {e}"
)
raise
embeddings = np.concatenate(embeddings_list)
# Explicit check (not assert): a count mismatch would silently
# truncate via zip() under `python -O`, mis-pairing vectors with
# docs. Raise instead so buffers stay intact for the next flush.
if len(embeddings) != len(docs_to_embed):
raise RuntimeError(
f"[{self.workspace}] Embedding count mismatch: expected "
f"{len(docs_to_embed)}, got {len(embeddings)}"
)
for i, ((_, pending_doc), embedding) in enumerate(
zip(docs_to_embed, embeddings), start=1
):
pending_doc.vector = embedding.tolist()
await _cooperative_yield(i)
# Deletes and upserts go as separate async_bulk requests so the
# delete record-count cap can differ from the upsert cap (mirrors
# mongo_impl's separate upsert/delete phases). The two buffers are
# disjoint -- delete() pops from pending_docs and upsert() discards
# from pending_deletes -- so request ordering is irrelevant.
delete_actions: list[dict[str, Any]] = [
{
"_op_type": "delete",
"_index": self._index_name,
"_id": doc_id,
}
for doc_id in pending_deletes
]
committed_doc_ids: set[str] = set()
index_actions: list[dict[str, Any]] = []
for doc_id, pending_doc in pending_docs.items():
if pending_doc.vector is None:
continue
committed_doc_ids.add(doc_id)
index_actions.append(
{
"_op_type": "index",
"_index": self._index_name,
"_id": doc_id,
"_source": {
**pending_doc.source,
"vector": pending_doc.vector,
},
}
)
if not delete_actions and not index_actions:
return
try:
# No per-operation refresh: search visibility is established
# by the refresh in index_done_callback().
log_prefix = f"[{self.workspace}] {self.namespace} flush:"
_, del_failed = await _run_chunked_async_bulk(
self.client,
delete_actions,
max_payload_bytes=self._max_upsert_payload_bytes,
max_records_per_batch=self._max_delete_records_per_batch,
log_prefix=log_prefix,
what="delete",
raise_on_error=False,
)
_, idx_failed = await _run_chunked_async_bulk(
self.client,
index_actions,
max_payload_bytes=self._max_upsert_payload_bytes,
max_records_per_batch=self._max_upsert_records_per_batch,
log_prefix=log_prefix,
what="upsert",
raise_on_error=False,
)
failed = list(del_failed) + list(idx_failed)
except OpenSearchException as e:
logger.error(
f"[{self.workspace}] Error flushing vector ops "
f"(upserts={len(pending_docs)}, "
f"deletes={len(pending_deletes)}): {e}"
)
# Bulk did not return per-doc statuses, so keep everything
# buffered for the next flush.
raise
retryable_ids, non_retryable_ops = _extract_bulk_failed_ids(failed)
# Keep ONLY retryable ops buffered for the next flush. Successful
# ops are popped; non-retryable (permanent 4xx) ops are dropped
# here, not retained: a permanently-unwritable op can never land,
# so keeping it would replay-and-refail on every later flush and
# poison every caller that shares this buffer — including direct
# flush paths that never run the pipeline's cleanup. The raise
# below (not retention) is what surfaces the failure and prevents
# a silent PROCESSED.
keep_ids = retryable_ids
for doc_id in committed_doc_ids:
if doc_id not in keep_ids:
pending_docs.pop(doc_id, None)
new_deletes: set[str] = {
doc_id for doc_id in pending_deletes if doc_id in keep_ids
}
pending_deletes.clear()
pending_deletes.update(new_deletes)
if retryable_ids:
logger.warning(
f"[{self.workspace}] {len(retryable_ids)} vector ops will "
f"retry on the next flush (transient failure)"
)
if non_retryable_ops:
sample = non_retryable_ops[:5]
sample_text = ", ".join(
f"{op.op}/{op.doc_id}/status={op.status}/{op.error}"
for op in sample
)
# A permanent (non-retryable) bulk failure means the data did
# not land. Raise so index_done_callback surfaces it and the
# pipeline aborts instead of marking the document PROCESSED.
raise RuntimeError(
f"[{self.workspace}] {self.namespace} flush: "
f"{len(non_retryable_ops)} vector ops failed permanently "
f"(non-retryable). Sample: {sample_text}"
)
async def query(
self, query: str, top_k: int, query_embedding: list[float] = None
) -> list[dict[str, Any]]:
"""k-NN similarity search with cosine score conversion for lucene engine."""
if not self._index_ready:
return []
if query_embedding is not None:
query_vector = (
query_embedding.tolist()
if hasattr(query_embedding, "tolist")
else list(query_embedding)
)
else:
embedding = await self.embedding_func(
[query], context="query", _priority=DEFAULT_QUERY_PRIORITY
)
query_vector = embedding[0].tolist()
search_body = {
"size": top_k,
"query": {"knn": {"vector": {"vector": query_vector, "k": top_k}}},
"_source": {"excludes": ["vector"]},
}
try:
response = await self.client.search(
index=self._index_name, body=search_body
)
results = []
for hit in response["hits"]["hits"]:
# OpenSearch k-NN with lucene engine and cosinesimil space type
# returns scores that can be used directly as similarity measure.
score = hit["_score"]
if score >= self.cosine_better_than_threshold:
doc = hit["_source"]
doc["id"] = hit["_id"]
doc["distance"] = score
results.append(doc)
logger.info(
f"[{self.workspace}] Vector query on {self._index_name}: "
f"top_k={top_k}, threshold={self.cosine_better_than_threshold}, "
f"total_hits={len(response['hits']['hits'])}, "
f"passed_filter={len(results)}, "
f"score_range=[{min((h['_score'] for h in response['hits']['hits']), default=0):.4f}, "
f"{max((h['_score'] for h in response['hits']['hits']), default=0):.4f}]"
)
return results
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return []
logger.error(f"[{self.workspace}] Error querying vectors: {e}")
return []
async def drop_pending_index_ops(self) -> None:
"""Discard buffered upserts/deletes (pipeline aborting on error)."""
async with self._flush_lock:
self._pending_vector_docs.clear()
self._pending_vector_deletes.clear()
async def index_done_callback(self) -> None:
"""Flush pending vector ops and refresh the index for k-NN visibility.
Flush runs first so that a previously-missing index gets recreated
by ``_flush_pending_vector_ops`` (via ``_ensure_index_ready``)
before any buffered writes are abandoned. The refresh step is
skipped only when the index is still not ready after the flush
attempt -- refreshing a half-built index is pointless.
Durability contract: each call embeds and bulk-indexes the *entire*
pending buffer in one shot. Deferred embedding runs inside
``_flush_pending_vector_ops``'s ``_flush_lock`` section (not in
``upsert``) precisely so this callback can guarantee every buffered
vector is embedded and flushed together; only transient per-doc
failures stay buffered for the next flush. Do not move embedding
out of the lock -- see ``_flush_pending_vector_ops`` for why.
"""
await self._flush_pending_vector_ops()
if not self._index_ready:
return
try:
await self.client.indices.refresh(index=self._index_name)
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return
raise
async def get_by_id(self, id: str) -> dict[str, Any] | None:
"""Get a vector document by ID, with read-your-writes against the buffer.
The ``vector`` field is stripped from the result to match every other
LightRAG vector backend (see ``NanoVectorDBStorage.get_by_id``).
Callers that need the embedding itself must use ``get_vectors_by_ids``.
"""
# Buffer lookups happen under the namespace lock so an in-flight
# flush is observed as either "completely before" or "completely
# after" -- never as a snapshot-swapped intermediate state.
async with self._flush_lock:
if id in self._pending_vector_deletes:
return None
pending = self._pending_vector_docs.get(id)
if pending is not None:
# pending.source is built in upsert from created_at + meta_fields
# and never carries the embedding, so no "vector" strip is needed
# here (unlike the mget path below, which excludes it server-side).
doc = dict(pending.source)
doc["id"] = id
return doc
if not self._index_ready:
return None
# Network IO outside the lock so mget RTT doesn't block flush.
try:
response = await _mget_optional_doc(
self.client,
self._index_name,
id,
source_excludes=["vector"],
)
if response is None:
return None
doc = response["_source"]
doc.pop("vector", None) # defensive in case _source_excludes is ignored
doc["id"] = response["_id"]
return doc
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return None
logger.error(f"[{self.workspace}] Error getting vector {id}: {e}")
return None
async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
"""Get multiple vector documents by IDs (read-your-writes), preserving order.
The ``vector`` field is stripped from each result; see ``get_by_id``.
"""
if not ids:
return []
buffered: dict[str, dict[str, Any] | None] = {}
remaining: list[str] = []
async with self._flush_lock:
for doc_id in ids:
if doc_id in self._pending_vector_deletes:
buffered[doc_id] = None
continue
pending = self._pending_vector_docs.get(doc_id)
if pending is not None:
# pending.source never carries the embedding; see get_by_id.
doc = dict(pending.source)
doc["id"] = doc_id
buffered[doc_id] = doc
continue
remaining.append(doc_id)
index_ready = self._index_ready
doc_map: dict[str, dict[str, Any] | None] = {}
if remaining and index_ready:
try:
response = await self.client.mget(
index=self._index_name,
body={"ids": remaining},
_source_excludes=["vector"],
)
for doc in response["docs"]:
if doc.get("found"):
data = doc["_source"]
data.pop("vector", None)
data["id"] = doc["_id"]
doc_map[doc["_id"]] = data
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
else:
logger.error(
f"[{self.workspace}] Error getting vectors by ids: {e}"
)
return [
buffered[doc_id] if doc_id in buffered else doc_map.get(doc_id)
for doc_id in ids
]
async def get_vectors_by_ids(self, ids: list[str]) -> dict[str, list[float]]:
"""Get vector embeddings for given IDs, with read-your-writes."""
if not ids:
return {}
result: dict[str, list[float]] = {}
remaining: list[str] = []
async with self._flush_lock:
docs_to_embed: list[tuple[str, _PendingVectorDoc]] = []
for doc_id in ids:
if doc_id in self._pending_vector_deletes:
continue
pending = self._pending_vector_docs.get(doc_id)
if pending is not None:
if pending.vector is None:
docs_to_embed.append((doc_id, pending))
else:
result[doc_id] = pending.vector
continue
remaining.append(doc_id)
index_ready = self._index_ready
if docs_to_embed:
contents = [pending_doc.content for _, pending_doc in docs_to_embed]
batches = [
contents[i : i + self._max_batch_size]
for i in range(0, len(contents), self._max_batch_size)
]
try:
embeddings_list = await asyncio.gather(
*[
self.embedding_func(batch, context="document")
for batch in batches
]
)
except Exception as e:
logger.error(
f"[{self.workspace}] Error lazily embedding pending vectors "
f"(upserts={len(docs_to_embed)}): {e}"
)
raise
embeddings = np.concatenate(embeddings_list)
# Explicit check (not assert): see _flush_pending_vector_ops.
if len(embeddings) != len(docs_to_embed):
raise RuntimeError(
f"[{self.workspace}] Embedding count mismatch: expected "
f"{len(docs_to_embed)}, got {len(embeddings)}"
)
for i, ((doc_id, pending_doc), embedding) in enumerate(
zip(docs_to_embed, embeddings), start=1
):
pending_doc.vector = embedding.tolist()
result[doc_id] = pending_doc.vector
await _cooperative_yield(i)
if not remaining:
return result
if not index_ready:
return result
try:
response = await self.client.mget(
index=self._index_name,
body={"ids": remaining},
_source_includes=["vector"],
)
for doc in response["docs"]:
if doc.get("found") and "vector" in doc.get("_source", {}):
result[doc["_id"]] = doc["_source"]["vector"]
return result
except OpenSearchException as e:
if _is_missing_index_error(e):
self._mark_index_missing()
return result
logger.error(f"[{self.workspace}] Error getting vectors: {e}")
return result
async def delete(self, ids: list[str]) -> None:
"""Buffer vector deletes for batched flush.
A delete cancels any pending upsert for the same id; the actual bulk
delete is performed by ``_flush_pending_vector_ops`` during the next
``index_done_callback`` / ``finalize`` call.
"""
if not ids:
return
if isinstance(ids, set):
ids = list(ids)
async with self._flush_lock:
for doc_id in ids:
self._pending_vector_docs.pop(doc_id, None)
self._pending_vector_deletes.add(doc_id)
logger.debug(
f"[{self.workspace}] Buffered delete for {len(ids)} vectors in {self.namespace}"
)
async def delete_entity(self, entity_name: str) -> None:
"""Buffer an entity vector delete by computing its hash ID."""
entity_id = compute_mdhash_id(entity_name, prefix="ent-")
async with self._flush_lock:
self._pending_vector_docs.pop(entity_id, None)
self._pending_vector_deletes.add(entity_id)
logger.debug(f"[{self.workspace}] Buffered delete for entity {entity_name}")
async def delete_entity_relation(self, entity_name: str) -> None:
"""Delete all relation vectors where entity appears as src or tgt.
The whole method runs under ``_flush_lock`` so the ``delete_by_query``
cannot interleave with an in-flight bulk indexing of a related doc.
Buffered upserts that match are pruned in-memory; persisted rows are
removed via the server-side ``delete_by_query``.
Buffer semantics — post-prune with caller short-circuit contract:
Matching pending upserts are pruned **only after** the
server-side ``delete_by_query`` succeeds (or returns the
equivalent of "index already missing"). On any other server
failure the exception is re-raised and the pending buffer
stays intact so a higher-level retry can still observe the
buffered relation vectors. Correctness relies on the caller
short-circuiting before ``index_done_callback`` can run;
``adelete_by_entity`` in ``utils_graph.py`` honors this.
Previously this method pre-pruned the buffer and swallowed
``OpenSearchException`` into a ``logger.error`` — that
combination silently dropped both the buffered relation
vectors and the server-side failure signal, leaving the
caller's graph + vector store permanently inconsistent.
"""
def _prune_pending() -> None:
for doc_id in [
k
for k, v in self._pending_vector_docs.items()
if v.source.get("src_id") == entity_name
or v.source.get("tgt_id") == entity_name
]:
self._pending_vector_docs.pop(doc_id, None)
async with self._flush_lock:
if not self._index_ready:
# No server state to mutate; buffer prune is the only
# delete intent we can record.
_prune_pending()
return
body = {
"query": {
"bool": {
"should": [
{"term": {"src_id": entity_name}},
{"term": {"tgt_id": entity_name}},
]
}
}
}
try:
# conflicts="proceed" tolerates stale search view after refresh removal.
await self.client.delete_by_query(
index=self._index_name, body=body, params={"conflicts": "proceed"}
)
except OpenSearchException as e:
if _is_missing_index_error(e):
# Index gone is equivalent to "all rows already
# deleted" — safe to prune pending and treat as
# success.
self._mark_index_missing()
_prune_pending()
return
logger.error(
f"[{self.workspace}] Error deleting relations for {entity_name}: {e}"
)
raise
# Server-side delete succeeded — safe to prune the pending
# buffer so subsequent flushes don't re-upsert the deleted
# relations.
_prune_pending()
logger.debug(
f"[{self.workspace}] Deleted relations for entity {entity_name}"
)
async def drop(self) -> dict[str, str]:
"""Delete and recreate the vector index, discarding pending buffers.
Runs entirely under ``_flush_lock`` so a concurrent flush / upsert
cannot land writes against an index that is being deleted and
rebuilt.
"""
async with self._flush_lock:
# Pending writes are meaningless once the index is dropped.
self._pending_vector_docs.clear()
self._pending_vector_deletes.clear()
try:
try:
await self.client.indices.delete(index=self._index_name)
logger.info(
f"[{self.workspace}] Dropped vector index: {self._index_name}"
)
except NotFoundError:
logger.info(
f"[{self.workspace}] Vector index already missing during drop: {self._index_name}"
)
# Recreate the index
await self._create_knn_index_if_not_exists()
self._index_ready = True
logger.info(
f"[{self.workspace}] Dropped and recreated vector index: {self._index_name}"
)
return {
"status": "success",
"message": f"Vector index {self._index_name} dropped and recreated",
}
except OpenSearchException as e:
self._mark_index_missing()
logger.error(f"[{self.workspace}] Error dropping vector index: {e}")
return {"status": "error", "message": str(e)}
except Exception as e:
self._mark_index_missing()
logger.error(
f"[{self.workspace}] Unexpected error dropping vector index: {e}"
)
return {"status": "error", "message": str(e)}
Reference in New Issue View Git Blame Copy Permalink