hkuds-lightrag-3056-3129

-rw-r--r-- 1 root root    3677 Nov  5 12:07 Dockerfile
-rw-r--r-- 1 root root    3603 Nov  5 12:07 Dockerfile.lite
-rw-r--r-- 1 root root    1070 Nov  5 12:07 LICENSE
-rw-r--r-- 1 root root     178 Nov  5 12:07 MANIFEST.in
-rw-r--r-- 1 root root    4782 Nov  5 12:07 Makefile
-rw-r--r-- 1
…[4144 chars omitted]…
namespace.py
./lightrag/operate.py
./lightrag/pipeline.py
./lightrag/prompt.py
./lightrag/prompt_multimodal.py
./lightrag/rerank.py
./lightrag/storage_migrations.py
./lightrag/table_markup.py
./lightrag/types.py
./lightrag/utils.py
./lightrag/utils_graph.py
./lightrag/utils_pipeline.py
./lightrag/_version.py
./lightrag/__init__.py
./lightrag.service.example
./lightrag_hku.egg-info/dependency_links.txt

lightrag/utils_graph.py:795:            await chunk_entity_relation_graph.upsert_edge(
lightrag/utils_graph.py:1102:            await chunk_entity_relation_graph.upsert_edge(
lightrag/utils_graph.py:1384:        await chunk_entity_relation_graph.upsert_edge(

        """
        Upsert a node in the Neo4j database.

        Args:
            node_id: The unique identifier for the node (used as label)
            node_data: Dictionary of node properties
        """
        if "entity_id" not in node_data:
            raise ValueError(
                "PostgreSQL: node properties must contain an 'entity_id' field"
            )

        # AGE supports binding scalar values in Cypher parameters here, but not
        # using a bound agtype obj
…[8167 chars omitted]…
ypher({_dollar_quote(self.graph_name)}, {_dollar_quote(cypher_query)}) AS (n agtype)"

        try:
            await self._query(query, readonly=False)
        except Exception as e:
            logger.error(f"[{self.workspace}] Error during node deletion: {e}")
            raise

    async def remove_nodes(self, node_ids: list[str]) -> None:
        """
        Remove multiple nodes from the graph.

2120:                graph_name=graph_name,
4993:def _is_transient_graph_write_error(exc: BaseException) -> bool:
5025:    def _get_workspace
…[5517 chars omitted]…
id = b.id
6600:                    SELECT start_id AS node_id FROM {self.graph_name}._ag_label_edge
6602:                    SELECT end_id AS node_id FROM {self.graph_name}._ag_label_edge
6611:                {self.graph_name}._ag_label_vertex v ON d.node_id = v.id
6646:                    {self.graph_name}._ag_label_vertex
6699:            drop_query = f"""SELECT * FROM cypher('{self.graph_name}', $$

def _is_transient_graph_write_error(exc: BaseException) -> bool:
    """Return True when a PGGraphQueryException wraps a transient write-time error.

    The inner _run_with_retry already handles connection-level transient errors
    (pool reset, TCP failures, etc.).  This predicate covers query-level transient
    errors that survive the connection layer and surface as PGGraphQueryException:
    deadlocks, serialization conflicts, and lock-acquisition timeouts that can
    occur under concurrent d
…[6338 chars omitted]…
e": e.__class__.__name__,
                }
            ) from e

        if data is None:
            result = []
        # decode records
        else:
            result = [self._record_to_dict(d) for d in data]

        return result

    async def has_node(self, node_id: str) -> bool:
        query = f"""
            SELECT EXISTS (
              SELECT 1
              FROM {self.graph_name}.base

tests/test_postgres_upsert_edge_cypher.py:108:async def test_upsert_edge_uses_parameterized_match_ids():
tests/test_postgres_upsert_edge_cypher.py:118:        await storage.upsert_edge("Node A", "Node B", {"weight": "1.0"})

from lightrag.kg.postgres_impl import PGGraphStorage


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def make_graph_storage() -> PGGraphStorage:
    """Construct a PGGraphStorage instance with a mocked _query method."""
    storage = PGGraphStorage.__new__(PGGraphStorage)
    storage.workspace = "test_ws"
    storage.namespace = "test_graph"
   
…[8629 chars omitted]…
:
        captured_execute_kwargs.append(kwargs)
        return None

    storage.db.execute = fake_execute

    test_params = {"params": json.dumps({"entity_id": "test"})}
    await storage._query(
        "SELECT 1",
        readonly=False,
        upsert=True,
        params=test_params,
    )

    assert len(captured_execute_kwargs) == 1
    assert captured_execute_kwargs[0]["data"] == test_params

310:            connection_params["ssl"] = ssl_context
315:                connection_params["ssl"] = True
317:                connection_params["ssl"] = False
331:                    connection_params["server_se
…[7102 chars omitted]…
      await connection.executemany(upsert_sql, batch_values)
3743:            await self.db._run_with_retry(_batch_upsert, timing_label=timing_label)
3755:    async def query(
3944:                        # Handle both pgvector-registered connections (returns list/tuple)
3945:                        # and non-registered connections (returns JSON string)
4930:            connection: asyncpg.Connection,

            operation: Async callable that receives an active connection.
            with_age: Whether to configure Apache AGE on the connection.
            graph_name: AGE graph name; required when with_age is True.

        Returns:
            The result returned by the operation.

        Raises:
            Exception: Propagates the last error if all retry attempts fail or a non-transient error occurs.
        """
        wait_strategy = (
            wait_exponential(
                multiplier=self.connection_retry_backoff,
    
…[8529 chars omitted]…
-wide PostgreSQL client pool shared by PG storages.

    The first successful initialization defines the pool configuration for the
    lifetime of the shared client. Reusing the pool with a different vector
    storage setup is not supported and will raise a fail-fast error.
    """

    _instances: dict[str, Any] = {
        "db": None,
        "ref_count": 0,
        "vector_signature": None,
    }

lightrag/chunker/paragraph_semantic.py:1086:    block dicts (no ``uuid`` / ``table_header`` propagation needed: the
lightrag/chunker/paragraph_semantic.py:1089:    if len(blocks) <= 1:
lightrag/chunker/paragraph_semantic.py:1090:        return blocks

                    direct_log(
                        f"== Lock == Process {self._pid}: Failed to release async lock after main lock failure: {inner_e}",
                        level="ERROR",
                        enable_output=True,
                    )

            raise

    def __enter__(self) -> "UnifiedLock[T]":
        """For backward compatibility"""

lightrag/lightrag.py:2902:                        async wi
…[3426 chars omitted]…

lightrag/utils_graph.py:190:    async with get_storage_keyed_lock(
lightrag/utils_graph.py:602:    async with get_storage_keyed_lock(
lightrag/utils_graph.py:767:    async with get_storage_keyed_lock(
lightrag/utils_graph.py:950:    async with get_storage_keyed_lock(
lightrag/utils_graph.py:1069:    async with get_storage_keyed_lock(
lightrag/utils_graph.py:1597:    async with get_storage_keyed_lock(

                source_id = chunk_to_source_map.get(source_chunk_id, "UNKNOWN")
                file_path = normalize_document_file_
…[7917 chars omitted]…
  prompt (Optional[str]): Custom prompts for fine-tuned control over the system's behavior. Defaults to None, which uses PROMPTS["rag_response"].

        Returns:
            str: The result of the query execution.
        """
        loop = always_get_an_event_loop()

        return loop.run_until_complete(self.aquery(query, param, system_prompt))  # type: ignore

    async def aquery(
        self,

        try:
            # Insert chunks into vector storage
            all_chunks_data: dict[str, dict[str, str]] = {}
            chunk_to_source_map: dict[str, str] = {}
            for chunk_data in custom_kg.get("chunks", []):
                chunk_content = sanitize_text_for_encoding(chunk_data["content"])
                source_id = chunk_data["source_id"]
                file_path = normalize_document_file_path(
                    chunk_data.get("file_path", "custom_kg")
                )
   
…[2539 chars omitted]…
. Please check the source mapping."
                    )

                node_data: dict[str, str] = {
                    "entity_id": entity_name,
                    "entity_type": entity_type,
                    "description": description,
                    "source_id": source_id,
                    "file_path": file_path,
                    "created_at": int(time.time()),
                }

    relation_chunks_storage=None,
) -> DeletionResult:
    """Asynchronously delete an entity and all its relationships.

    Also cleans up entity_chunks_storage and relation_chunks_storage to remove chunk tracking.

    Args:
        chunk_entity_relation_graph: Graph storage instance
        entities_vdb: Vector database storage for entities
        relationships_vdb: Vector database storage for relationships
        entity_name: Name of the entity to delete
        entity_chunks_storage: Optional KV storage for tracking chunks that reference this entity
        relati
…[8629 chars omitted]…
aise RuntimeError(
            "Shared data not initialized. Call initialize_share_data() before using locks!"
        )
    async_lock = _async_locks.get("data_init_lock") if _is_multiprocess else None
    return UnifiedLock(
        lock=_data_init_lock,
        is_async=not _is_multiprocess,
        name="data_init_lock",
        enable_logging=enable_logging,
        async_lock=async_lock,
    )

import tempfile
import pytest
import numpy as np
from unittest.mock import AsyncMock

from lightrag.kg.networkx_impl import NetworkXStorage
from lightrag.kg.shared_storage import initialize_share_data
from lightrag.utils import EmbeddingFunc, make_relation_vdb_ids


# ---------------------------------------------------------------------------
# Helpers
# ------------------------------------------------
…[8630 chars omitted]…
tor_data"] == {"ok": True}
            assert [
                call.args[0] for call in rag.relationships_vdb.get_by_id.await_args_list
            ] == [
                normalized_rel_id,
                legacy_rel_id,
                normalized_rel_id,
                legacy_rel_id,
            ]

            await rag.finalize_storages()


class TestPostgresBatchOrdering:
    @pytest.mark.offline

            "chunks": [
                {
                    "content": "chunk content",
                    "chunk_order_index": 0,
                    "source_id": "src-1",
                }
            ],
            "entities": [
                {
                    "entity_name": "EntityA",
                    "entity_type": "CONCEPT",
                    "description": "An entity",
                    "source_id": "src-1",
                    "file_path": "test.pdf",
                }
            ],
       
…[8629 chars omitted]…
wargs["ordered"] is True

    @pytest.mark.offline
    @pytest.mark.asyncio
    async def test_upsert_edges_batch_deduplicates_source_node_upserts(self):
        pytest.importorskip("pymongo")
        from lightrag.kg.mongo_impl import MongoGraphStorage

        storage = MongoGraphStorage.__new__(MongoGraphStorage)
        storage.collection = AsyncMock()
        storage.edge_collection = AsyncMock()

        {
            "host": "localhost",
            "port": 5432,
            "user": "postgres",
            "password": "postgres",
            "database": "postgres",
            "workspace": "test_ws",
            "max_connections": 10,
            "connection_retry_attempts": 3,
            "connection_retry_backoff": 0,
            "connection_retry_backoff_max": 0,
            "pool_close_timeout": 5.0,
        }
    )


@pytest.mark.asyncio
async def test_execute_timing_logs_success():
    db = make_db()

    async def fake_run_with_retry(operation, **kwargs):
        conn = AsyncMock()
        conn
…[2702 chars omitted]…
def test_performance_timing_logs_ignores_old_env(monkeypatch):
    with monkeypatch.context() as m:
        m.setenv("LIGHTRAG_DOC_QUERY_TIMING_LOGS", "true")
        m.setenv("LIGHTRAG_PERFORMANCE_TIMING_LOGS", "false")
        reloaded = importlib.reload(utils_module)
        assert reloaded.PERFORMANCE_TIMING_LOGS is False

    importlib.reload(utils_module)
tests/test_description_api_validation.py

        batch_build_start = time.perf_counter()
        for i, (k, v) in enumerate(data.items(), start=1):
            try:
                batch.append(
                    (
                        self.workspace,
                        k,
                        v["content_summary"],
                        v["content_
…[2928 chars omitted]…

            drop_sql = SQL_TEMPLATES["drop_specifiy_table_workspace"].format(
                table_name=table_name
            )
            await self.db.execute(drop_sql, {"workspace": self.workspace})
            return {"status": "success", "message": "data dropped"}
        except Exception as e:
            return {"status": "error", "message": str(e)}


class PGGraphQueryException(Exception):

        Args:
            None

        Returns:
            str: The graph name for the current workspace
        """
        workspace = self.workspace
        namespace = self.namespace

        if workspace and workspace.strip() and workspace.strip().lower() != "default":
            # Ensure names comply with PostgreSQL identifier specifications
            safe_workspace = re.sub(r"[^a-zA-
…[8629 chars omitted]…
                      json_content, f"single {type_id} for field {k}"
                        )
                        d[k] = parsed_data if parsed_data is not None else None
                    else:
                        # May be other types of agtype data, keep as is
                        d[k] = v
            else:
                d[k] = v  # Keep as string

        return d

    @staticmethod

    httpx = None
from dataclasses import InitVar, asdict, dataclass, field, replace
from datetime import datetime, timezone
from functools import partial
from typing import (
    Any,
    AsyncIterator,
    Awaitable,
    Callable,
    Iterator,
    cast,
    final,
    Literal,
    Mapping,
    Optional,
    List,
    Dict,
    Union,
)
from lightrag.prompt import (
    PROMPTS,
    get_default_entity_extraction_prompt_profile,
    resolve_entity_extraction_prompt_profile,
    validate_entity_extraction_prompt_profile_for_mode,
)
from lightrag.constants import (
    DEFAULT_CHUNK_P_SIZE,
    DEFAULT_MAX_GLEANING,
    DEFAULT_MAX_EXTRACTION_RECORDS,
    DEFAULT_MAX_EXTRACTI
…[2936 chars omitted]…
118:                set_default_workspace(self.workspace)
1119:            elif default_workspace != self.workspace:
1121:                    f"Creating LightRAG instance with workspace='{self.workspace}' "
1128:            await initialize_pipeline_status(workspace=self.workspace)
2760:            "pipeline_status", workspace=self.workspace
2763:            "pipeline_status", workspace=self.workspace

1039:            workspace=self.workspace,
1118:                set_default_workspace(self.workspace)
1119:            elif default_workspace != self.workspace:
1121:                    f"Creating LightRAG instance with workspace='{self.workspace}' "
1128:            await initialize_pipeline_status(workspace=self.workspace)
2760:            "p
…[5936 chars omitted]…
errors.

    Args:
        namespace: The namespace to get the lock for.
        workspace: Workspace identifier (may be empty string for global namespace)
        enable_logging: Whether to enable lock operation logging

    Returns:
        NamespaceLock: A reusable lock wrapper that can be used with 'async with'

    Example:
        lock = get_namespace_lock("pipeline_status", workspace="space1")

def _dollar_quote(s: str, tag_prefix: str = "AGE") -> str:
    """
    Generate a PostgreSQL dollar-quoted string with a unique tag.

    PostgreSQL dollar-quoting uses $tag$ as delimiters. If the content contains
    the same delimiter (e.g., $$ or $AGE1$), it will break the query.
    This function finds a unique tag that doesn't conflict with the content.

    Args:
        s: The string to quote
        tag_prefix: Prefix for generating unique tags (default: "AGE")

    Returns:
        The dollar-quote
…[951 chars omitted]…
t$AGE1$
126-
127-    Example:
128-        >>> _dollar_quote("hello")
129-        '$AGE1$hello$AGE1$'
130-        >>> _dollar_quote("$AGE1$ test")
131-        '$AGE2$$AGE1$ test$AGE2$'
132-        >>> _dollar_quote("$$$")  # Content with dollar signs
133-        '$AGE1$$$$AGE1$'
134-    """
135-    s = "" if s is None else str(s)
136-    for i in itertools.count(1):
137-        tag = f"{tag_prefix}{i}"

tests/test_workspace_migration_isolation.py:270:        async def mock_run_with_retry(operation, *args, **kwargs):
tests/test_workspace_migration_isolation.py:275:        db._run_with_retry = AsyncMock(side_effect=mock_run_with_retry)

import pytest
import numpy as np
from unittest.mock import AsyncMock, MagicMock
from lightrag.kg.postgres_impl import PGDocStatusStorage, PGKVStorage, PGVectorStorage
from lightrag.namespace import NameSpace
from lightrag.utils import EmbeddingFunc


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

GLOBAL_CONFIG = {"embedding_batch_num": 10}


def make_storage(namespace
…[6495 chars omitted]…

        await storage.upsert({"k": {"v": 1}})


# ---------------------------------------------------------------------------
# 11. Multiple records go into one batch when within limit
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_multiple_records_single_batch():
    storage = make_storage(NameSpace.KV_STORE_FULL_DOCS)
    data = {

 .github/workflows/copilot-setup-steps.yml          |   58 +
 .github/workflows/docker-build-lite.yml            |  113 +
 .github/workflows/docker-build-manual.yml          |  109 +
 .github/workflows/docker-publish.yml               |  120 +
 .github/workflows/li
…[547 chars omitted]…
                                 |  111 +
 LICENSE                                            |   21 +
 MANIFEST.in                                        |    5 +
 Makefile                                           |   99 +
 README-zh.md                                       |  436 ++
 README.assets/b2aaf634151b4706892693ffb43d9093.png |  Bin 0 -> 482231 bytes
becebcc HEAD@{0}: commit (initial): base

        # Vector configuration
        _ev = config.get("enable_vector", True)
        self.enable_vector = (
            _ev
            if isinstance(_ev, bool)
            else str(_ev).lower() in ("true", "1", "yes", "on")
        )  # True for backward compatibility, can be set to False to disable vector features
        self.vector_index_type = config.get("vector_index_type")
        self.hnsw_m = config.get("hnsw_m")
        self.hnsw_ef = config.get("hns
…[4002 chars omitted]…
lse None
        logger.warning(
            "PostgreSQL transient connection issue on attempt %s/%s: %r",
            retry_state.attempt_number,
            self.connection_retry_attempts,
            exc,
        )
        await self._reset_pool()

    async def _run_with_retry(
        self,
        operation: Callable[[asyncpg.Connection], Awaitable[T]],
        *,
        with_age: bool = False,

541-        self._default_enable_logging = default_enable_logging
542-        self._async_lock: Dict[str, asyncio.Lock] = {}  # local keyed locks
543-        self._async_lock_count: Dict[
544-            str, int
545-        ] = {}  # local keyed locks referenced count
546-        self._async_lock_cleanup_data: Dict[
547-            str, time.time
548-        ] = {} 
…[8629 chars omitted]…
tely create and add entry to list (critical for rollback to work)
858-                    entry = {
859-                        "key": key,
860-                        "lock": lock,
861-                        "entered": False,
862-                        "debug_inc": False,
863-                        "ref_incremented": True,  # Mark that reference count has been incremented
864-                    }

            self.message = exception["message"] if "message" in exception else "unknown"
            self.details = exception["details"] if "details" in exception else "unknown"
        else:
            self.message = exception
            self.details = "unknown"

    def get_message(self) -> str:
        return self.message

    def get_details(self) -> Any:
        return self.details


def _is_transient_graph_write_error(exc: BaseException) -> boo
…[306 chars omitted]…
e as PGGraphQueryException:
    deadlocks, serialization conflicts, and lock-acquisition timeouts that can
    occur under concurrent document ingestion.
    """
    if not isinstance(exc, PGGraphQueryException):
        return False
    cause = exc.__cause__
    if cause is None:
        return False
    return isinstance(
        cause,
        (
            asyncpg.exceptions.DeadlockDetectedError,

        self,
        query: str,
        readonly: bool = True,
        upsert: bool = False,
        params: dict[str, Any] | None = None,
        timing_label: str | None = None,
        lock_statement: tuple[str, tuple[Any, ...]] | None = None,
    ) -> list[dict[str, Any]]:
        """
        Query the graph by taking a cypher query, converting it to an
        age compatible query, executing it and converting the result

        Args:
            query (str): a cypher query to be 
…[7487 chars omitted]…
    readonly=False,
                upsert=True,
                params=pg_params,
                timing_label=timing_label,
            )
            performance_timing_log(
                "[%s] total complete in %.4fs source_node_id=%s target_node_id=%s",
                timing_label,
                time.perf_counter() - total_start,
                source_node_id,
                target_node_id,

5797:            await self._query(query, readonly=False)
5820:                await self._query(query, readonly=False)
5881:            results = await self._query(query, params={"ids
…[1642 chars omitted]…
es: List of (node_id, node_data) tuples.
        """
        if not nodes:
            return
        deduped_nodes: dict[str, dict[str, str]] = {}
        for node_id, node_data in nodes:
            deduped_nodes.pop(node_id, None)
            deduped_nodes[node_id] = node_data

        for node_id, node_data in deduped_nodes.items():
            await self.upsert_node(node_id, node_data=node_data)

                {"entity_id": node_id},
                ensure_ascii=False,
            )
        }
        timing_label = f"{self.workspace} PGGraphStorage.upsert_node"
        total_start = time.perf_counter()
        performance_timing_log(
            "[%s] start node_id=%s",
            timing_label,
            node_id,
        )

        try:
            await self._query(
                query,
                readonly=False,
                upsert=True,
        
…[171 chars omitted]…
 source_node_id, target_node_id),
                ),
            )
            performance_timing_log(
                "[%s] total complete in %.4fs node_id=%s",
                timing_label,
                time.perf_counter() - total_start,
                node_id,
            )

        except Exception:
            performance_timing_log(
                "[%s] total failed after %.4fs node_id=%s",

                upsert=True,
                params=pg_params,
                timing_label=timing_label,
            )
            performance_timing_log(
                "[%s] total complete in %.4fs node_id=%s",
                timing_label,
                time.perf_counter() - total_start,
                node_id,
            )

        except Exception:
            performance_timing_log(
                "[%s] total failed after %.4fs node_id=%s",
                timing_label,
                time.perf_counter() - total_start,
        try:
            await self._query(
                query,
                r
…[312 chars omitted]…
e_timing_log(
                "[%s] total complete in %.4fs source_node_id=%s target_node_id=%s",
                timing_label,
                time.perf_counter() - total_start,
                source_node_id,
                target_node_id,
            )

        except Exception:
            performance_timing_log(
                "[%s] total failed after %.4fs source_node_id=%s target_node_id=%s",

    BaseGraphStorage,
    BaseKVStorage,
    BaseVectorStorage,
    DocProcessingStatus,
    DocStatus,

    async def ainsert_custom_kg(
        self,
        custom_kg: dict[str, Any],
        full_doc_id: str = None,
    ) -> None:
        # Coordinate custom-KG writes with the entity/relationship keyed locks
        # used by the regular graph-ingestion pipeline.  Acquire every name
        # this batch may write up front so overlapping batches cannot race on
        # shared entities or endpoints.
        graph_lock_keys: set[str] = set()
        for entity_data in custom_kg.get("entities", []) or []:
            entity_name = entity_data.get("entity_name")
     
…[2896 chars omitted]…
ert(data_for_rels_vdb),
            )

            if legacy_rel_ids_to_delete:
                await self.relationships_vdb.delete(legacy_rel_ids_to_delete)

        except Exception as e:
            logger.error(f"Error in ainsert_custom_kg: {e}")
            raise
        finally:
            try:
                if update_storage:
                    await self._insert_done()
            finally:

            resolved_timeout = TimeoutSauce(connect=timeout, read=timeout)
    
        try:
            resp = conn.urlopen(
                method=request.method,
                url=url,
                body=request.body,  # type: ignore[arg-type]  # urllib3 stubs don't accept Iterable[bytes | str]
                headers=request.headers,  # type: ignore[arg-type]  
…[2109 chars omitted]…
py::TestAinsertCustomKgBatchPath::test_ainsert_custom_kg_deduplicates_entities_and_undirected_edges
FAILED tests/test_batch_graph_operations.py::TestAinsertCustomKgBatchPath::test_ainsert_custom_kg_keeps_legacy_relation_rows_if_upsert_fails
FAILED tests/test_batch_graph_operations.py::TestAinsertCustomKgBatchPath::test_get_relation_info_falls_back_to_legacy_relation_vdb_id
6 failed, 41 passed in 3.56s

    raise new_e
/usr/local/lib/python3.12/site-packages/urllib3/connectionpool.py:470: in _make_request
 
…[4071 chars omitted]…
ionError("HTTPSConnection(host='openaipublic.blob.core.windows.net', port=443): Failed to resolve 'openaipublic.blob.core.windows.net' ([Errno -3] Temporary failure in name resolution)"))
=========================== short test summary info ============================
FAILED tests/test_batch_graph_operations.py::TestAinsertCustomKgBatchPath::test_ainsert_custom_kg_calls_batch_methods
1 failed in 1.29s

  File "<stdin>", line 24, in main
  File "/usr/local/lib/python3.12/site-packages/tenacity/asyncio/__init__.py", line 193, in async_wrapped
    return await copy(fn, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/tenacity/asyncio/__init__.py", line 112, in __c
…[1111 chars omitted]…
ostgres_impl.py", line 5669, in upsert_edge
    await self._query(
  File "/repo/lightrag/kg/postgres_impl.py", line 5407, in _query
    result = [self._record_to_dict(d) for d in data]
              ^^^^^^^^^^^^^^^^^^^^^^^
  File "/repo/lightrag/kg/postgres_impl.py", line 5218, in _record_to_dict
    for k in record.keys():
             ^^^^^^^^^^^
AttributeError: 'str' object has no attribute 'keys'

    return future.result()
           ^^^^^^^^^^^^^^^
  File "<stdin>", line 27, in main
AssertionError

10 | | from typing import Any, Awaitable, Callable, TypeVar, Union, final
11 | | import numpy as np
12 | | import configparser
13 | | import ssl
14 | | import itertools
15 | |
16 | | from lightrag.types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge
17 | |
18 | | from tenacity import (
19 | |     AsyncRetrying,
20 | |     RetryCallState,
21 | |     retry,
22 | |     retry_if_exception,
23 | |     retry_if_exception_type,
24 | |     stop_after_attempt,
25 | |     wait_exponential,
26 | |     wait_fixed,
27 | | )
28 | |
29 | | from ..base import (
30 | |     
…[760 chars omitted]…
      list_data
     -         ), f"Embedding count mismatch: expected {len(list_data)}, got {len(embeddings)}"
3691 +         assert len(embeddings) == len(list_data), (
3692 +             f"Embedding count mismatch: expected {len(list_data)}, got {len(embeddings)}"
3693 +         )
3694 |         embedding_fill_start = time.perf_counter()
     |

1 file would be reformatted, 1 file already formatted

     ) -> list[dict[str, Any]]:
         """
         Query the graph by taking a cypher query, converting it to an
@@ -5318,12 +5319,47 @@ class PGGraphStorage(BaseGraphStorage):
                 (e.g. ``{"params": json.dumps({"entity_id": "..."})}``).
                 Honoured f
…[6499 chars omitted]…
_kg: {e}")
             raise
         finally:
-            if update_storage:
-                await self._insert_done()
+            try:
+                if update_storage:
+                    await self._insert_done()
+            finally:
+                if graph_lock_context is not None:
+                    await graph_lock_context.__aexit__(None, None, None)
 
     def query(
         self,

208-        relationships_vdb=DummyVectorStorage(),
209-        entity_name="EntityA",
210-        updated_data={"entity_type": "ORG"},
211-    )
212-
213-    assert result["operation_summary"]["operation_status"] == "success"
214-
215-
216-def test_handle_single_relationship_extraction_ignores_empty_description():
217-    relation = _handle_single_relationship_extraction(
218-        ["relation", "Alice", "Bob", "works_with", "   "],
219-        chunk_key="chunk-1",
220-        timestamp=1,
221-    )

        except MaxRetryError as e:
            if isinstance(e.reason, ConnectTimeoutError):
                # TODO: Remove this in 3.0.0: see #2811
                if not isinstance(e.reason, NewConnectionError):
                    raise ConnectTimeout(e, request=request)
    
            if isinstance(e.reason, ResponseError):
                raise RetryError(e, request=request)
    
            if isinstance(e.reason, _ProxyError):
                raise ProxyError(e, request=request)
    
     
…[1197 chars omitted]…
py::TestAinsertCustomKgBatchPath::test_ainsert_custom_kg_deduplicates_entities_and_undirected_edges
FAILED tests/test_batch_graph_operations.py::TestAinsertCustomKgBatchPath::test_ainsert_custom_kg_keeps_legacy_relation_rows_if_upsert_fails
FAILED tests/test_batch_graph_operations.py::TestAinsertCustomKgBatchPath::test_get_relation_info_falls_back_to_legacy_relation_vdb_id
6 failed, 80 passed in 3.92s

        self._log_llm_role_config("initialized")

        self._storages_status = StoragesStatus.CREATED

    async def initialize_storages(self):
        """Storage initialization must be called one by one to prevent deadlock"""
        if self._storages_status == StoragesStatus.CREATED:
            # Set the first initialized workspace will set the default workspace
            # Allows namespace operation without specifying workspace for backward compatibility
…[398 chars omitted]…
"
                )

            # Auto-initialize pipeline_status for this workspace
            from lightrag.kg.shared_storage import initialize_pipeline_status

            await initialize_pipeline_status(workspace=self.workspace)

            for storage in (
                self.full_docs,
                self.text_chunks,
                self.full_entities,
                self.full_relations,

+    asyncpg.exceptions.QueryCanceledError,
+)
+
+
+def _is_transient_graph_query_error(exc: BaseException) -> bool:
+    """Return True for query-level transient errors that should be retried."""
+    return isinstance(exc, _TRANSIENT_GRAPH_QUERY_ERRORS)
+
+
 def _is_transient_graph_write_error(exc: BaseException) -> bool:
     """Return True when a PGGraphQueryException wraps a transient write-time error.
 
@@ -5004,15 +5017,7 @@ def _is_transient_g
…[6959 chars omitted]…
                  {
+                        "message": f"Error executing graph query: {query}",
+                        "wrapped": query,
+                        "detail": repr(e),
+                        "error_type": e.__class__.__name__,
+                    }
+                ) from e
             raise
 
     async def upsert_nodes_batch(self, nodes: list[tuple[str, dict[str, str]]]) -> None:

@@ -1449,6 +1450,32 @@ class LightRAG(_RoleLLMMixin, _StorageMigrationMixin, _PipelineMixin):
         custom_kg: dict[str, Any],
         full_doc_id: str = None,
     ) -> None:
+        # Coordinate custom-KG writes with the entity/relationship keyed locks
+        # used by the regular graph-ingestion pipeline.  Acquire every name
+        # this batch may write up front so overlapping batches cannot race on
+        # shared entities or endpoints.
+        graph_lock_keys: set[str] = set()
+        for entity_data in custom_kg.get("entities",
…[1132 chars omitted]…
_kg: {e}")
             raise
         finally:
-            if update_storage:
-                await self._insert_done()
+            try:
+                if update_storage:
+                    await self._insert_done()
+            finally:
+                if graph_lock_context is not None:
+                    await graph_lock_context.__aexit__(None, None, None)
 
     def query(
         self,

from dotenv import load_dotenv

# use the .env that is inside the current folder
# allows to use different .env file for each lightrag instance
# the OS environment variables take precedence over the .env file
load_dotenv(dotenv_path=".env", override=False)

T = TypeVar("T")

# PostgreSQL identifier length limit (in bytes)
PG_MAX_IDENTIFIER_LENGTH = 63

# All known vector index suffixes, used to drop conflicting indexes when switching types
_VECTOR_INDEX_SUFFIXES = [
    "hnsw_cosine",
    "hnsw_halfvec_cosine",
    "ivfflat_cosine",
    "vchordrq_cosine",
]


def _safe_index_name(table_name: str, index_suffix: str) -> str:
    """
    Generate a PostgreSQL-safe index name that won't be truncated.

    PostgreSQL silently truncates identifiers to 63 bytes. This function
    ensures index names stay within that limit by hashing long table names.

-rw-r--r-- 1 root root      41 Nov  2 19:28 CLAUDE.md
-rw-r--r-- 1 root root    3677 Nov  2 19:28 Dockerfile
-rw-r--r-- 1 root root    3603 Nov  2 19:28 Dockerfile.lite
-rw-r--r-- 1 root root    1070 Nov  2 19:28 LICENSE
-rw-r--r-- 1 root root     178 Nov  2 19:28 MANIFEST.in
-rw-r--r-- 1 root root    4782 Nov  2 19:28 Makefile
-rw-r-
…[4137 chars omitted]…
namespace.py
./lightrag/operate.py
./lightrag/pipeline.py
./lightrag/prompt.py
./lightrag/prompt_multimodal.py
./lightrag/rerank.py
./lightrag/storage_migrations.py
./lightrag/table_markup.py
./lightrag/types.py
./lightrag/utils.py
./lightrag/utils_graph.py
./lightrag/utils_pipeline.py
./lightrag/_version.py
./lightrag/__init__.py
./lightrag.service.example
./lightrag_hku.egg-info/dependency_links.txt

lightrag/kg/networkx_impl.py
lightrag/kg/opensearch_impl.py
lightrag/kg/postgres_impl.py
lightrag/kg/qdrant_impl.py
lightrag/kg/redis_impl.py
lightrag/kg/shared_storage.py
lightrag/kg/__init__.py
lightrag/tools/lightrag_visualizer/graph_visualizer.py
lightrag/utils_graph.py
tests/test_batch_graph_operations.py
tests/test_graph_storage.py
tests/test_memgraph_storage.py
tests/test_paragraph_semantic_merge_and_fallback.py
tests/test_paragraph_semantic_split_long_block.py
tests/test_paragraph_semantic_table_split.py
tests/tes
…[2927 chars omitted]…
await chunk_entity_relation_graph.upsert_edge(
lightrag/utils_graph.py:338:                        await chunk_entity_relation_graph.upsert_edge(
lightrag/utils_graph.py:795:            await chunk_entity_relation_graph.upsert_edge(
lightrag/utils_graph.py:1102:            await chunk_entity_relation_graph.upsert_edge(
lightrag/utils_graph.py:1384:        await chunk_entity_relation_graph.upsert_edge(

import ssl
import itertools

from lightrag.types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge

from tenacity import (
    AsyncRetrying,
    RetryCallState,
    retry,
    retry_if_exception,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
    wait_fixed,
)

from ..base import (
    BaseGraphStorage,
    BaseKVStorage,
    BaseVectorStorage,
    DocProcessingStatus,
    DocStatus,
    DocStatusStorage,
)
from ..exceptions import DataMigrationError
from ..namespace import NameSpace, is_namespace
from ..utils import logger, _cooperative_yield, performance_timing_log
from ..kg.share
…[8630 chars omitted]…
TACH DELETE n"""

        query = f"SELECT * FROM cypher({_dollar_quote(self.graph_name)}, {_dollar_quote(cypher_query)}) AS (n agtype)"

        try:
            await self._query(query, readonly=False)
        except Exception as e:
            logger.error(f"[{self.workspace}] Error during node removal: {e}")
            raise

    async def remove_edges(self, edges: list[tuple[str, str]]) -> None:

                        )
                    else:
                        logger.warning(
                            "PostgreSQL, SSL client certificate or key file not found"
                        )

                # Load certificate revocation list if provided
    
…[8632 chars omitted]…
o ensure atomic graph and vector db operations
lightrag/utils_graph.py:1069:    async with get_storage_keyed_lock(
lightrag/utils_graph.py:1204:    It should only be called by public APIs that have already acquired necessary locks.
lightrag/utils_graph.py:1221:        Caller must acquire appropriate locks before calling this function.
lightrag/utils_graph.py:1597:    async with get_storage_keyed_lock(

import pytest
import numpy as np
from unittest.mock import AsyncMock, MagicMock
from lightrag.kg.postgres_impl import PGDocStatusStorage, PGKVStorage, PGVectorStorage
from lightrag.namespace import NameSpace
from lightrag.utils import EmbeddingFunc


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

GLOBAL_CONFIG = {"embedding_batch_num": 10}


def make_storage(namespace
…[8629 chars omitted]…
n []

    with patch.object(storage, "_query", side_effect=fake_query):
        await storage.upsert_edge("Node A", "Node B", {"weight": "1.0"})

    call = captured_calls[0]
    sql = call["sql"]
    assert "entity_id: $src_id" in sql
    assert "entity_id: $tgt_id" in sql
    params = json.loads(call["params"]["params"])
    assert params["src_id"] == "Node A"
    assert params["tgt_id"] == "Node B"

            entity_nodes: list[tuple[str, dict[str, str]]] = []
            for entity_data in deduped_entities.values():
                entity_name = entity_data["entity_name"]
                entity_type = entity_data.get("entity_type", "UNKNOWN")
                description = entity_data.get
…[4926 chars omitted]…
normalized_tgt_id,
                        "description": relationship_data["description"],
                        "keywords": relationship_data["keywords"],
                        "source_id": source_id,
                        "weight": relationship_data.get("weight", 1.0),
                        "file_path": file_path,
                        "created_at": int(time.time()),
                    }

lightrag/lightrag.py:1445:        loop.run_until_complete(self.ainsert_custom_kg(custom_kg, full_doc_id))
lightrag/lightrag.py:1447:    async def ainsert_custom_kg(
lightrag/lightrag.py:1672:            logger.error(f"Error in ainsert_custom_kg: {e}")

            ConnectionError,
            OSError,
            asyncpg.exceptions.InterfaceError,
            asyncpg.exceptions.TooManyConnectionsError,
            asyncpg.exceptions.CannotConnectNowError,
            asyncpg.exceptions.PostgresConnectionError,
            asyncpg.exceptions.ConnectionDoesNotExistError,
            asyncpg.exceptions.ConnectionFailureError,
        )

 
…[8629 chars omitted]…
PACE environment variable: '{self.db.workspace}' (overriding '{self.workspace}/{self.namespace}')"
                )
                self.workspace = self.db.workspace
            elif hasattr(self, "workspace") and self.workspace:
                # Use storage class's workspace (medium priority)
                pass
            else:
                # Use "default" for compatibility (lowest priority)

    httpx = None
from dataclasses import InitVar, asdict, dataclass, field, replace
from datetime import datetime, timezone
from functools import partial
from typing import (
    Any,
    AsyncIterator,
    Awaitable,
    Callable,
    Iterator,
    cast,
    final,
    Literal,
    Mapping,
    Optional,
    List,
    Dict,
    Union,
)
from lightrag.prompt import (
    PROMPTS,
    get_default_entity_extraction_prompt_profile,
    resolve_entity_extraction_prompt_profile,
    validate_entity_extraction_prompt_profile_for_mode,
)
from lightrag.constants import (
    DEFAULT_CHUNK_P_SIZE,
    DEFAULT_MAX_GLEANING,
    DEFAULT_MAX_EXTRACTION_RECORDS,
    DEFAULT_MAX_EXTRACTI
…[3648 chars omitted]…
le default workspace is set to '{default_workspace}'"
1125:            # Auto-initialize pipeline_status for this workspace
1126:            from lightrag.kg.shared_storage import initialize_pipeline_status
1128:            await initialize_pipeline_status(workspace=self.workspace)
2760:            "pipeline_status", workspace=self.workspace
2763:            "pipeline_status", workspace=self.workspace

            json_content = agtype_str[:last_double_colon]
            type_identifier = agtype_str[last_double_colon + 2 :]

            return json_content, type_identifier

        @staticmethod
        def safe_json_parse(json_str: str, context: str = "") -> dict:
            """
            Safe JSON parsing with simplified error logging
            """
            try:
                return json.loads(json_str)
            except json
…[8630 chars omitted]…
th"] == "test.pdf"

            await rag.finalize_storages()

    @pytest.mark.offline
    @pytest.mark.asyncio
    async def test_ainsert_custom_kg_no_hasattr_needed(self):
        """
        The batch methods are always available on the base class, so no
        hasattr() guard should be needed. Verify that a storage backend
        implementing only the abstract methods (no batch overrides) still

def _require_non_empty_description(
    description: Any, *, operation: str, object_type: str
) -> None:
    if description is None or not str(description).strip():
        raise ValueError(
            f"{object_type.capitalize()} description cannot be empty for {operation} operation"
        )


async def _persist_graph_updates(
    entities_vdb=None,
    relationships_vdb=None,
    chunk_entity_relation_graph=None,
    entity_chunks_storage=None,
    relation_chunks_storage=None,
) -> None:
    """Unified callback to persist updates after graph operations.

    En
…[8630 chars omitted]…
H_FIELD_SEP
            - "join_unique_comma": Join unique items separated by comma and space
            - "max": Keep the maximum numeric value (for numeric fields)
        filter_none_only: If True, only filter None values (keep empty strings, 0, etc.).
            If False, filter all falsy values. Default is False for backward compatibility.

    Returns:
        Dictionary containing merged data

        captured_calls.append({"sql": sql, **kwargs})
        return []

    with patch.object(storage, "_query", side_effect=fake_query):
        await storage.upsert_edge(
            "Alice", "Bob", {"weight": "1.0", "description": "knows"}
        )

    assert len(captured_calls) == 1
    call = captured_calls[0]
    assert "$1::agtype" in call["sql"]
    assert '"Alice"' not in call["sql"].replace("$1::agtype", "")
    assert '"Bob"' no
…[4789 chars omitted]…
ync def test_upsert_edge_handles_empty_props():
tests/test_postgres_upsert_edge_cypher.py:101:        await storage.upsert_edge("Alice", "Bob", {})
tests/test_postgres_upsert_edge_cypher.py:108:async def test_upsert_edge_uses_parameterized_match_ids():
tests/test_postgres_upsert_edge_cypher.py:118:        await storage.upsert_edge("Node A", "Node B", {"weight": "1.0"})
5271:    def _format_properties(

829-        # to avoid deadlocks
830-        self._keys = sorted(keys)
831-        self._enable_logging = (
832-            enable_logging
833-            if enable_logging is not None
834-            else parent._default_enable_logging
835-        )
836-        self._ul: Optional[List[Dict[str, Any]]] = None  # set in __aenter__
837-
838-    # ----- enter -----
839-    async def __aenter__(self):
840-        if self._ul is not None:
841-            raise RuntimeError("KeyedUnifiedLock already acquired in current cont
…[6404 chars omitted]…
ongoGraphStorage.__new__(MongoGraphStorage)
        storage.collection = AsyncMock()

        await MongoGraphStorage.upsert_nodes_batch(
            storage,
            [
                ("EntityA", _make_node("EntityA")),
82:from lightrag.kg.shared_storage import (
818:        from lightrag.kg.shared_storage import (
1126:            from lightrag.kg.shared_storage import initialize_pipeline_status

    """Storage backend for key-value data."""

    vector_storage: str = field(default="NanoVectorDBStorage")
    """Storage backend for vector embeddings."""

    graph_storage: str = field(default="NetworkXStorage")
    """Storage backend for knowledge graphs."""

    doc_status_storage: str = field(default="JsonDocStatusStorage")
    """Storage type for tracking document processing statuses."""

    # Workspace
    # ---

    workspace: str = field(default_factory=lambda: os.getenv("WORKSPACE", ""))
    """Workspace for data isolation. Defaults to empty string if WORKSPACE environment variable 
…[8628 chars omitted]…
ce)

            for storage in (
                self.full_docs,
                self.text_chunks,
                self.full_entities,
                self.full_relations,
                self.entity_chunks,
                self.relation_chunks,
                self.entities_vdb,
                self.relationships_vdb,
                self.chunks_vdb,
                self.chunk_entity_relation_graph,

        async with pipeline_status_lock:
            if pipeline_status.get("cancellation_requested", False):
                raise PipelineCancelledException("User cancelled during merge phase")

    # Collect all nodes and edges from all chunks
    all_nodes = defaultdict(list)
    all_edges = defaultdic
…[8629 chars omitted]…
,
2930-                        entities,
2931-                        knowledge_graph_inst,
2932-                        entity_vdb,
2933-                        global_config,
778:    def _build_global_config(self) -> dict[str, Any]:
925:        global_config = self._build_global_config()
1865:        global_config = self._build_global_config()
1983:        global_config = self._build_global_config()

    def __init__(self, *, default_enable_logging: bool = True) -> None:
        self._default_enable_logging = default_enable_logging
        self._async_lock: Dict[str, asyncio.Lock] = {}  # local keyed locks
        self._async_lock_count: Dict[
            str, int
        ] = {}  # local keyed locks referenced count
        self._async_lock_cleanup_data: Dict[
            str, time.time
        ] = {}  # local keyed locks timeout
        self._mp_lock
…[2451 chars omitted]…
) is not None:
            raise SystemExit(
                "ERROR: ENTITY_TYPES environment variable is no longer supported. "
                "Please customize entity type guidance through the prompt template instead. "
                "Set addon_params={'entity_types_guidance': '...'} or replace the prompt template."
            )

        self._replace_addon_params(addon_params, mark_dirty=False)

tests/test_postgres_upsert_edge_cypher.py:100:    with patch.object(storage, "_query", side_effect=fake_query):
tests/test_postgres_upsert_edge_cypher.py:117:    with patch.object(storage, "_query", side_effect=fake_query):

"""

import json
import pytest
from unittest.mock import MagicMock, patch

from lightrag.kg.postgres_impl import PGGraphStorage


# ---------------------------------------------------------------------------
# Helpers
# -------------------------------------------------------------------------
…[8629 chars omitted]…
n []

    with patch.object(storage, "_query", side_effect=fake_query):
        await storage.upsert_edge("Node A", "Node B", {"weight": "1.0"})

    call = captured_calls[0]
    sql = call["sql"]
    assert "entity_id: $src_id" in sql
    assert "entity_id: $tgt_id" in sql
    params = json.loads(call["params"]["params"])
    assert params["src_id"] == "Node A"
    assert params["tgt_id"] == "Node B"

- **scripts/**: `test.sh` (preferred test runner), `setup/` interactive environment wizard (use `make env-*` rather than calling `setup.sh` directly — see *Configuration > Setup Wizard Outputs*), and release tooling.
- **tests/** and root-level `test_*.py`: Pytest coverage. Working datasets stay in `inputs/`,
…[8629 chars omitted]…
ding model you MUST clear the data directory (optionally keeping `kv_store_llm_response_cache.json` for LLM cache). Existing vectors will not match the new model's space.

### Storage Configuration

Configure via environment variables or constructor params:

```python
# Environment-based (recommended for production)
queries: 1
distinct after delete: False
contains lock: False
empty props pattern: True

                # Use storage class's workspace (medium priority)
                pass
            else:
                # Use "default" for compatibility (lowest priority)
                self.workspace = "default"

            # Dynamically generate
…[3295 chars omitted]…
one

    async def index_done_callback(self) -> None:
        # PG handles persistence automatically
        pass

    @staticmethod
    def _record_to_dict(record: asyncpg.Record) -> dict[str, Any]:
        """
        Convert a record returned from an age query to a dictionary
679:                f"select create_graph('{graph_name}')"
5118:                f"SELECT create_graph('{self.graph_name}')",

./requirements-offline-storage.txt
./tests/test_graph_storage.py
./tests/test_json_doc_status_storage.py
./tests/test_llm_binding_image_payload.py
./tests/test_memgraph_storage.py
./tests/test_mongo_storage.py
./tests/test_opensearch_storage.py
./tests/test_postgres_age_quote_fix.py
./tests/test_postgres_client_manager.py
./tests/test_vlm_image_inputs.py
./tests/tes
…[264 chars omitted]…
           outgoing_query = f"SELECT * FROM cypher({_dollar_quote(self.graph_name)}::name, {_dollar_quote(outgoing_cypher)}::cstring, $1::agtype) AS (node_id text, connected_id text)"
./lightrag/kg/postgres_impl.py:6131:            incoming_query = f"SELECT * FROM cypher({_dollar_quote(self.graph_name)}::name, {_dollar_quote(incoming_cypher)}::cstring, $1::agtype) AS (node_id text, connected_id text)"

        This method:
        - Sets the PostgreSQL `search_path` to include `ag_catalog`, ensuring that Apache AGE functions can be used without specifying the schema.
        - Attempts to create a new gra
…[8628 chars omitted]…
 await self._query(
                query,
                readonly=False,
                upsert=True,
                params=pg_params,
                timing_label=timing_label,
                transaction_lock=(lock_query, lock_params),
            )
            performance_timing_log(
                "[%s] total complete in %.4fs source_node_id=%s target_node_id=%s",
                timing_label,

except Exception:  # pragma: no cover - optional dependency
    httpx = None
from dataclasses import InitVar, asdict, dataclass, field, replace
from datetime import datetime, timezone
from functools import partial
from typing import (
    Any,
    AsyncIterator,
    Awaitable,
    Callable,
    Iterator,
    cast,
    final,
    Literal,
    Mapping,
    Optional,
    List,
    Dict,
    Union,
)
from lightrag.prompt import (
    PROMPTS,
    get_default_entity_extraction_prompt_profile,
    resolve_entity_extraction_prompt_profile,
    validate_entity_extraction_prompt_profile_for_mode,
)
from lightrag.constants import (
    DEFAULT_CHUNK_P_SIZE,
    DEFAULT_MAX_GLEANING,
    DEFAULT_MAX_EXTRACTIO
…[8629 chars omitted]…
           "file_path": dp.get("file_path", "custom_kg"),
                    }
                    for dp in all_relationships_data
                }

                legacy_rel_ids_to_delete = sorted(
                    {
                        rel_id
                        for dp in all_relationships_data
                        for rel_id in make_relation_vdb_ids(dp["src_id"], dp["tgt_id"])[1:]

 lightrag/kg/postgres_impl.py |  97 +++++++++----
 lightrag/lightrag.py         | 333 ++++++++++
…[8629 chars omitted]…
ues():
-                src_id = relationship_data["src_id"]
-                tgt_id = relationship_data["tgt_id"]
-                source_chunk_id = relationship_data.get("source_id", "UNKNOWN")
-                source_id = chunk_to_source_map.get(source_chunk_id, "UNKNOWN")
-                file_path = normalize_document_file_path(
-                    relationship_data.get("file_path", "custom_kg")

try:
    import httpx
except Exception:  # pragma: no cover - optional dependency
    httpx = None
from dataclasses import InitVar, asdict, dataclass, field, replace
from datetime import datetime, timezone
from functools import partial
from typing import (
    Any,
    AsyncIterator,
    Awaitable,
    Callable,
    Iterator,
    cast,
    final,
    Literal,
    Mapping,
    Optional,
    List,
    Dict,
    Union,
)
from lightrag.prompt import (
    PROMPTS,
    get_default_entity_extraction_prompt_profile,
    resolve_entity_extraction_prompt_profile,
    validate_entity_extraction_prompt_profile_for_mode,
)
from lightrag.constants import (
    DEFAULT_CHUNK_P_S
…[8629 chars omitted]…
terator[str]:
        """
        Perform a async query (backward compatibility wrapper).

        This function is now a wrapper around aquery_llm that maintains backward compatibility
        by returning only the LLM response content in the original format.

        Args:
            query (str): The query to be executed.
            param (QueryParam): Configuration parameters for query execution.

            return []
    
        with patch.object(storage, "_query", side_effect=fake_query):
            await storage.upsert_edge("Alice", "Bob", {"weight": "0.5"})
    
        sql = captured_sql[0]
        assert "OPTIONAL MATCH (source)-[old:DIRECTED]-(target)" in sql
        assert "DELETE old" in sql
        assert "CREATE (sou
…[8631 chars omitted]…
_undirected_edges
FAILED tests/test_batch_graph_operations.py::TestAinsertCustomKgBatchPath::test_ainsert_custom_kg_keeps_legacy_relation_rows_if_upsert_fails
FAILED tests/test_batch_graph_operations.py::TestAinsertCustomKgBatchPath::test_get_relation_info_falls_back_to_legacy_relation_vdb_id
!!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 10 failures !!!!!!!!!!!!!!!!!!!!!!!!!!
10 failed, 32 passed in 3.63s

>           sock = connection.create_connection(
                (self._dns_host, self.port),
                self.timeout,
                source_address=self.source_address,
                socket_options=self.socket_options,
            )

/usr/local/lib/python3.12/site-packages/urllib3/connection.py:239: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/usr/local/lib/python3.12/site-packages/urllib3/util/conne
…[8630 chars omitted]…
 in name resolution)"))

/usr/local/lib/python3.12/site-packages/requests/adapters.py:729: ConnectionError
=========================== short test summary info ============================
FAILED tests/test_batch_graph_operations.py::TestAinsertCustomKgBatchPath::test_ainsert_custom_kg_calls_batch_methods
!!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 1 failures !!!!!!!!!!!!!!!!!!!!!!!!!!!
1 failed in 1.30s

        transaction_lock: tuple[str, tuple[Any, ...]] | None = None,
    ) -> list[dict[str, Any]]:
        """
        Query the graph by taking a cypher query, converting it to an
        age compatible query, executing it and converting the result

        Args:
            query (str): a cypher query to be executed
            readonly (bool): if True, uses db.query; if False, uses db.execute.
                Both paths support the ``params`` argument.
            upsert (bool): passed through to db.execute for write operations.
            params (d
…[8628 chars omitted]…
 historical serial fallback.

        Args:
            nodes: List of (node_id, node_data) tuples.
        """
        if not nodes:
            return
        deduped_nodes: dict[str, dict[str, str]] = {}
        for node_id, node_data in nodes:
            deduped_nodes.pop(node_id, None)
            deduped_nodes[node_id] = node_data

71	26	lightrag/kg/postgres_impl.py
180	153	lightrag/lightrag.py

                            if prepared_params:
                                await connection.execute(query, *prepared_params)
                            else:
                                await connection.execute(qu
…[299 chars omitted]…
tement_start,
                                )

                    data = await self.db._run_with_retry(
                        _execute_transaction,
                        with_age=True,
                        graph_name=self.graph_name,
                        timing_label=timing_label,
                    )
                if timing_label:
                    performance_timing_log(

                direct_log(
                    f"Lock rollback unexpected error for {entry['key']}: {e}",
                    level="ERROR",
                    enable_output=True,
                )

        self._ul = None

    # ----- exit -----
    async def __aexit__(self, exc_type, exc, tb):
        if self._ul is None:
            return

        async def release_all_locks():
            """Release all locks with comprehensive error handling, protected fro
…[5520 chars omitted]…
d in make_relation_vdb_ids(dp["src_id"], dp["tgt_id"])[1:]
                    }
                )

                # Parallel VDB upserts (was serial in original)
                await asyncio.gather(
                    self.entities_vdb.upsert(data_for_entities_vdb),
                    self.relationships_vdb.upsert(data_for_rels_vdb),
                )

                if legacy_rel_ids_to_delete:

@@ -84,6 +85,7 @@ from lightrag.kg.shared_storage import (
     get_default_workspace,
     set_default_workspace,
     get_namespace_lock,
+    get_storage_keyed_lock,
 )
 
 from lightrag.base import (
@@ -1489,184 +1491,209 @@ class LightRAG(_RoleLLMMixin, _StorageMigrationMixin, _PipelineMixin):
                     self.text_chunks.upsert(all_chunks_data),
                 )
 
-            # Keep the last declaration for each entity_name so batch backends
-            # preserve the old serial upsert semantics deterministically.
-            deduped_entities: dict[str, dict[str, An
…[8629 chars omitted]…
path": dp.get("file_path", "custom_kg"),
+                    }
                     for dp in all_relationships_data
-                    for rel_id in make_relation_vdb_ids(dp["src_id"], dp["tgt_id"])[1:]
                 }
-            )
 
-            # Parallel VDB upserts (was serial in original)
-            await asyncio.gather(
-                self.entities_vdb.upsert(data_for_entities_vdb),

                self.full_relations,
                self.entity_chunks,
                self.relation_chunks,
                self.llm_response_cache,
                self.entities_vdb,
                self.relationships_vdb,
                self.chunks_vdb,
                self.chunk_entity_relation_graph,
            ]
1683:class TiktokenTokenizer(Tokenizer):
        # Init Embedding
        # Step 1: Capture embedding_func and max_token_size before applying rate_limit decorator
        or
…[2606 chars omitted]…
ment status storage
        self.doc_status_storage_cls = get_storage_class(self.doc_status_storage)

        self.llm_response_cache: BaseKVStorage = self.key_string_value_json_storage_cls(  # type: ignore
            namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
            workspace=self.workspace,
            global_config=global_config,
            embedding_func=self.embedding_func,
        )

        Raises:
            ImportError: If tiktoken is not installed.
            ValueError: If the model_name is invalid.
        """
        try:
            import tiktoken
        except ImportError:
            raise ImportError(
                "tiktoken is not installed. Please install it with `pip install tiktoken` or define custom `tokenizer_func`."
            )

        try:
            tokenizer = tiktoken.encoding_for_model(model_name)
            super().__init__(model_name=model_name, tokenizer=tokenizer)
…[1044 chars omitted]…
rovided parameters
        if self.tokenizer is None:
            if self.tiktoken_model_name:
                self.tokenizer = TiktokenTokenizer(self.tiktoken_model_name)
            else:
                self.tokenizer = TiktokenTokenizer()

        # Initialize ollama_server_infos if not provided
        if self.ollama_server_infos is None:
            self.ollama_server_infos = OllamaServerInfos()

        }
        property_entries = []
        for index, (key, value) in enumerate((edge_data or {}).items()):
            parameter_name = f"edge_prop_{index}"
            safe_key = str(key).replace("
…[2127 chars omitted]…
     transaction_lock=(lock_query, lock_params),
            )
            performance_timing_log(
                "[%s] total complete in %.4fs source_node_id=%s target_node_id=%s",
                timing_label,
                time.perf_counter() - total_start,
                source_node_id,
                target_node_id,
            )

        except Exception:
            performance_timing_log(

=========================== short test summary info ============================
FAILED tests/test_postgres_upsert_edge_cypher.py::test_upsert_edge_contains_optional_match_delete_create
FAILED tests/test_postgres_cypher_injection.py::test_upsert_edge_uses_parameterized_cypher
FAILED tests/test_postgres_cypher_injection.py::test_upsert_edge_unicode_entity_ids
3 failed, 35 passed in 1.92s

1623-
1624-
1625:class Tokenizer:
1626-    """
1627-    A wrapper around a tokenizer to provide a consistent interface for encoding and decoding.
1628-    """
1629-
1630-    def __init__(self, model_name: str, tokenizer: TokenizerInterface):
1631-        """
1632-        Initializes the Tokenizer with a tokenizer model name and a tokenizer instance.
1633-
1634-        Args:
1635-            model_name: The associated model name for the tokenizer.
1636-            tokeni
…[7748 chars omitted]…
th the base default

            async def base_upsert_nodes_batch(self, nodes):
                for node_id, node_data in nodes:
                    await self.upsert_node(node_id, node_data=node_data)

            storage.upsert_node = spy  # type: ignore[assignment]
            await base_upsert_nodes_batch(storage, nodes)

            assert call_log == ["NodeA", "NodeB"]

    @pytest.mark.offline

            source_node_id (str): Label of the source node (used as identifier)
            target_node_id (str): Label of the target node (used as identifier)
            edge_data (dict): dictionary of properties to set on the edge
        """
        # AGE requires edge property keys to be present in the CREATE map, but
        # accepts bound Cypher parameters for the corresponding values. Keeping
        # those values out of the query text prevents a payloa
…[1854 chars omitted]…
LEAST($2::text, $3::text) || "
            "E'\\x01' || GREATEST($2::text, $3::text), 0))"
        )
        lock_params = (self.graph_name, source_node_id, target_node_id)
        timing_label = f"{self.workspace} PGGraphStorage.upsert_edge"
        total_start = time.perf_counter()
        performance_timing_log(
            "[%s] start source_node_id=%s target_node_id=%s",
            timing_label,

unformatted: File would be reformatted
    --> lightrag/lightrag.py:1533:51
     |
1532 |                     entity_type = entity_data.get("entity_type", "UNKNOWN")
     -                     description = entity_data.get("description", "No description provided")
1533 +                     description = entity_data.get(
1534 +       
…[1076 chars omitted]…
lationships_data
     -                         for rel_id in make_relation_vdb_ids(dp["src_id"], dp["tgt_id"])[1:]
1691 +                         for rel_id in make_relation_vdb_ids(dp["src_id"], dp["tgt_id"])[
1692 +                             1:
1693 +                         ]
1694 |                     }
     |

2 files would be reformatted
 M lightrag/kg/postgres_impl.py
 M lightrag/lightrag.py

        try:
            # Insert chunks into vector storage
            all_chunks_data: dict[str, dict[str, str]] = {}
            chunk_to_source_map: dict[str, str] = {}
            for chunk_data in custom_kg.get("chunks", []):
                chunk_content = sanitize_text_for_encoding(chunk_data["content"])
                source_id = chunk_data["source_id"]
                file_path = normalize_document_file_path(
                    chunk_data.get("file_path", "custom_kg")
                )
  
…[5188 chars omitted]…
None,
    ) -> str | Iterator[str]:
        """
        Perform a sync query.

        Args:
            query (str): The query to be executed.
            param (QueryParam): Configuration parameters for query execution.
 lightrag/kg/postgres_impl.py | 119 ++++++++++++++++++++++++++++++++-----------
 lightrag/lightrag.py         |  35 ++++++++++++-
 2 files changed, 123 insertions(+), 31 deletions(-)

                    self.text_chunks.upsert(all_chunks_data),
                )

            # Hold the same workspace-scoped entity locks used by regular graph
            # ingestion while this batch checks and updates graph data. Include
            # relationship endpoints even when they are not declared as entities:
            # the batch may create placeholder nodes for them.
            graph_lock_keys = {
…[8628 chars omitted]…
s_to_delete:
                await self.relationships_vdb.delete(legacy_rel_ids_to_delete)

        except Exception as e:
            logger.error(f"Error in ainsert_custom_kg: {e}")
            raise
        finally:
            try:
                if graph_lock_acquired and graph_lock_context is not None:
                    await graph_lock_context.__aexit__(None, None, None)
            finally:

  1542	                        f"Entity '{entity_name}' has an UNKNOWN source_id. Please check the source mapping."
  1543	                    )
  1544	
  1545	      
…[1310 chars omitted]…
        str: The result of the query execution.
  1725	        """
try:
    import httpx
except Exception:  # pragma: no cover - optional dependency
    httpx = None
from dataclasses import InitVar, asdict, dataclass, field, replace
from datetime import datetime, timezone
from functools import partial
from typing import (
    Any,
    AsyncIterator,
    Awaitable,
    Callable,
    Iterator,
    cast,

                sorted_edge_key,
                namespace=namespace,
                enable_logging=False,
            ):
                try:
                    added_entities = []  # Track entities added during edge processing

                    edge_data = await _merge_edges_then_upsert(
             
…[6724 chars omitted]…
get_storage_keyed_lock(
1011:            async with get_storage_keyed_lock(
2887:            sorted_edge_key = tuple(sorted(edge_key))
2888:            all_edges[sorted_edge_key].extend(edges)
2923:            async with get_storage_keyed_lock(
3028:            sorted_edge_key = sorted([edge_key[0], edge_key[1]])
3031:            async with get_storage_keyed_lock(
3032:                sorted_edge_key,

        }
        property_entries = []
        for index, (key, value) in enumerate((edge_data or {}).items()):
            parameter_name = f"edge_prop_{index}"
            safe_key = str(key).replace("
…[574 chars omitted]…
   CREATE (source)-[r:DIRECTED {props_literal}]->(target)
                     RETURN r"""

        # Graph name and all Cypher values are query parameters. Keep the AGE
        # parameter first so the SQL remains independent of caller data.
        query = (
            f"SELECT * FROM cypher($2::text::name, "
            f"{_dollar_quote(cypher_query)}::cstring, $1::agtype) AS (r agtype)"
        )

            "graph_name": self.graph_name,
        }
        # Advisory lock keys use the workspace graph and an order-independent
        # endpoint pair. Positional SQL parameters ensure the values are never
        # part of the statement text.
        lock_query = (
            "SELECT pg_advisory_xact_lock(hashtextextended("
            "$1::text || E'\\x01' || LEAST($2::text, $3::text) || "
            "E'\\x01' || GREATEST($2::text, $3::text), 0))"

        cypher_query = f"""MATCH (source:base {{entity_id: $src_id}})
                     WITH source
                     MATCH (target:base {{entity_id: $tgt_id}})
                     WITH source, target
       
…[1362 chars omitted]…
 f"{self.workspace} PGGraphStorage.upsert_edge"
        total_start = time.perf_counter()
        performance_timing_log(
            "[%s] start source_node_id=%s target_node_id=%s",
            timing_label,
            source_node_id,
            target_node_id,
        )

        try:
            await self._query(
                query,
                readonly=False,
                upsert=True,

311:class TestAinsertCustomKgBatchPath:
350:    async def test_ainsert_custom_kg_calls_batch_methods(self):
388:    async def test_ainsert_custom_kg_canonicalizes_file_pa
…[479 chars omitted]…
 async def test_upsert_nodes_batch_preserves_last_write_wins(self):
789:    async def test_upsert_edges_batch_preserves_last_write_wins(self):
812:class TestMongoBatchOrdering:
815:    async def test_upsert_nodes_batch_uses_ordered_bulk_write(self):
834:    async def test_upsert_edges_batch_uses_ordered_bulk_write(self):
854:    async def test_upsert_edges_batch_deduplicates_source_node_upserts(self):

@@ -84,6 +85,7 @@ from lightrag.kg.shared_storage import (
     get_default_workspace,
     set_default_workspace,
     get_namespace_lock,
+    get_storage_keyed_lock,
 )
 
 from lightrag.base import (
@@ -1331,6 +1333,8 @@ class LightRAG(_RoleLLMMixin, _StorageMigrationMixin, _PipelineMixin):
         self, full_text: str, text_chunks: list[str], doc_id: str | None = None
     ) -> None:
         update_storage = False
+        graph_lock_context = None
+        graph_lock_acquired = False
         try:
             # Clean input texts
             full_text = sanitize_text_for_enco
…[1923 chars omitted]…
is not None:
+                    await graph_lock_context.__aexit__(*sys.exc_info())
+            finally:
+                if update_storage:
+                    await self._insert_done()
 
     def query(
         self,
 lightrag/kg/postgres_impl.py | 151 +++++++++++++++++++++++++++++--------------
 lightrag/lightrag.py         |  36 ++++++++++-
 2 files changed, 135 insertions(+), 52 deletions(-)

    async def ainsert_custom_chunks(
        self, full_text: str, text_chunks: list[str], doc_id: str | None = None
    ) -> None:
        update_storage = False
        graph_lock_context = None
        graph_lock_acquired = False
        try:
            # Clean input texts
            full_text = sanitize_text_for_encoding(full_text)
            text_chunks = [sanitize_text_for_encoding(chunk) for chunk in text_chunks]
            file_path = normalize_document_file_path("")
        if pipeline_status is not None and pipelin
…[368 chars omitted]…
p.run_until_complete(self.ainsert_custom_kg(custom_kg, full_doc_id))

    async def ainsert_custom_kg(
        self,
        custom_kg: dict[str, Any],
        full_doc_id: str = None,
    ) -> None:
        update_storage = False
        try:
            # Insert chunks into vector storage
            all_chunks_data: dict[str, dict[str, str]] = {}
            chunk_to_source_map: dict[str, str] = {}

@@ -84,6 +85,7 @@ from lightrag.kg.shared_storage import (
     get_default_workspace,
     set_default_workspace,
     get_namespace_lock,
+    get_storage_keyed_lock,
 )
 
 from lightrag.base import (
@@ -1450,6 +1452,8 @@ class LightRAG(_RoleLLMMixin, _StorageMigrationMixin, _PipelineMixin):
         full_doc_id: str = None,
     ) -> None:
         update_storage = False
+        graph_lock_context = None
+        graph_lock_acquired = False
         try:
             # Insert chunks into vector storage
             all_chunks_data: dict[str, dict[str, str]] = {}
@@ -1489,6 +1493,
…[2030 chars omitted]…
ge:
+                    await self._insert_done()
 
     def query(
         self,
1455:        graph_lock_context = None
1512:                graph_lock_context = get_storage_keyed_lock(
1517:                await graph_lock_context.__aenter__()
1704:                if graph_lock_acquired and graph_lock_context is not None:
1705:                    await graph_lock_context.__aexit__(*sys.exc_info())

            return []
    
        with patch.object(storage, "_query", side_effect=fake_query):
            await storage.upsert_edge("Alice", "Bob", {"weight": "0.5"})
    
        sql = captured_sql[0]
        assert "OPTIONAL MATCH (source)-[old:DIRECTED]-(target)" in sql
        assert "DELETE ol
…[3243 chars omitted]…
AssertionError
=========================== short test summary info ============================
FAILED tests/test_postgres_upsert_edge_cypher.py::test_upsert_edge_contains_optional_match_delete_create
FAILED tests/test_postgres_cypher_injection.py::test_upsert_edge_uses_parameterized_cypher
FAILED tests/test_postgres_cypher_injection.py::test_upsert_edge_unicode_entity_ids
3 failed, 60 passed in 1.93s

198-        self.pool_close_timeout = config["pool_close_timeout"]
199-        logger.info(
200-            "PostgreSQL, Retry config: attempts=%s, backoff=%.1fs, backoff_max=%.1fs, pool_close_timeout=%.1fs",
201-            self.connection_retry_attempts,
202-            self.connection_retry_backoff,
203-            self.connection_re
…[3407 chars omitted]…
elf._transient_exceptions),
442-                wait=wait_strategy,
443-                before_sleep=self._before_sleep,
--
527-        async for attempt in AsyncRetrying(
528-            stop=stop_after_attempt(self.connection_retry_attempts),
529:            retry=retry_if_exception_type(self._transient_exceptions),
530-            wait=wait_strategy,
531-            before_sleep=self._before_sleep,

DEBUG_LOCKS = False


# Define a direct print function for critical logs that must be visible in all processes
def direct_log(message, enable_output: bool = True, level: str = "DEBUG"):
    """
    Log a message directly to stderr to ensure visibility in all processes,
    including the Gunicorn master process.

    Args:
        message: The message to log
        level: Log level for message (control the visibility of the message by comparing with the current logger level)
        enable_output: Enable or disable log message (Force to
…[7058 chars omitted]…
y_guard, \
1202-        _internal_lock, \
1203-        _data_init_lock, \
1204-        _shared_dicts, \
1205-        _init_flags, \
1206-        _initialized, \
1207-        _update_flags, \
1208-        _async_locks, \
1209-        _storage_keyed_lock, \
1210-        _earliest_mp_cleanup_time, \
1211-        _last_mp_cleanup_time
1212-
1213-    # Check if already initialized
1214-    if _initialized:

        _internal_lock, \
        _data_init_lock, \
        _shared_dicts, \
        _init_flags, \
        _initialized, \
        _update_flags, \
        _async_locks, \
        _storage_keyed_lock, \
        _earliest_mp_cleanup_time, \
        _last_mp_cleanup_time

    # Check if already initialized
    if _initialized:
        direct_log(
            f"Process {os.getpid()} Shared-Data already initialized (multiprocess={_is_multiprocess})"
        )
        return

    _workers = workers

    if workers > 1
…[1768 chars omitted]…
    If None or empty string, uses the default workspace set by
                   set_default_workspace().
    """
    pipeline_namespace = await get_namespace_data(
        "pipeline_status", first_init=True, workspace=workspace
    )

    async with get_internal_lock():
        # Check if already initialized by checking for required fields
        if "busy" in pipeline_namespace:
            return

+        )
         embedding_fill_start = time.perf_counter()
         for i, d in enumerate(list_data, start=1):
             d["__vector__"] = embeddings[i - 1]

5322-        age compatible query, executing it and converting the result
5323-
5324-        Args:
5325-            query (str): a cypher query to be executed
5326-            readonly (bool): if True, uses db.query; if False, uses db.execute.
5327-                Both paths support the ``params`` argument.
5328-            upsert (bool): passed through to db.execute for write operations.
5329-            params (dict | None): AGE agtype paramete
…[6657 chars omitted]…
 result
5428-
5429-    async def has_node(self, node_id: str) -> bool:
5430-        query = f"""
5431-            SELECT EXISTS (
5432-              SELECT 1
5433-              FROM {self.graph_name}.base
5434-              WHERE ag_catalog.agtype_access_operator(
5435-                      VARIADIC ARRAY[properties, '"entity_id"'::agtype]
5436-                    ) = (to_json($1::text)::text)::agtype

        {
            "host": "localhost",
            "port": 5432,
            "user": "postgres",
            "password": "postgres",
            "database": "postgres",
            "workspace": "test_ws",
            "max_connections": 10,
            "connection_retry_attempts": 3,
            "connection_retry_backoff": 0,
            "connection_retry_backoff_max": 0,
            "pool_close_timeout": 5.0,
        }
    )


@pytest.mark.asyncio
async def test_execute_timing_logs_success():
    db = make_db()

    async def fake_run_with_retry(operation, **kwargs):
        conn = AsyncMock()
        conn
…[2709 chars omitted]…
reSQLDB(db_config)
        attempt_count = {"value": 0}

        # Original create_pool function
        original_create_pool = asyncpg.create_pool

        async def mock_create_pool_with_failures(*args, **kwargs):
            """Mock that fails first 2 times, then calls real create_pool."""
            attempt_count["value"] += 1
            print(f"  → Connection attempt {attempt_count['value']}")

            return []
    
        with patch.object(storage, "_query", side_effect=fake_query):
            await storage.upsert_edge("Alice", "Bob", {"weight": "0.5"})
    
        sql = captured_sql[0]
        assert "OPTIONAL MATCH (source)-[old:DIRECTED]-(target)" in sql
        assert "DELETE ol
…[4177 chars omitted]…
========
FAILED tests/test_postgres_upsert_edge_cypher.py::test_upsert_edge_contains_optional_match_delete_create
FAILED tests/test_postgres_upsert_edge_cypher.py::test_upsert_edge_handles_empty_props
FAILED tests/test_postgres_cypher_injection.py::test_upsert_edge_uses_parameterized_cypher
FAILED tests/test_postgres_cypher_injection.py::test_upsert_edge_unicode_entity_ids
4 failed, 59 passed in 1.71s

            logger.info(
                f"[{self.workspace}] PostgreSQL Graph initialized: graph_name='{self.graph_name}'"
            )

            # Create AGE extension and configure graph environment once at initialization
            # Use _run_with_retry so transient connection errors are retried and pool=None
            # is handled safely (unlike a bare pool.acquire
…[5571 chars omitted]…
 configuration
        self.connection_retry_attempts = config["connection_retry_attempts"]
        self.connection_retry_backoff = config["connection_retry_backoff"]
        self.connection_retry_backoff_max = max(
            self.connection_retry_backoff,
            config["connection_retry_backoff_max"],
        )
        self.pool_close_timeout = config["pool_close_timeout"]
        logger.info(

                    direct_log(
                        f"Lock release reference error for key {key}: {e}",
                        level="ERROR",
                        enable_output=True,
                    )

                # 3. Decrement debug counter
                if debug_inc:
                    try:
                        dec_debug_n_locks_acquired()
              
…[4528 chars omitted]…
f.entities_vdb,
                self.relationships_vdb,
                self.chunks_vdb,
                self.chunk_entity_relation_graph,
            ]
            if storage_inst is not None
        ]
        await asyncio.gather(*tasks)

        log_message = "In memory DB persist to disk"
        logger.info(log_message)

        if pipeline_status is not None and pipeline_status_lock is not None:

@@ -84,6 +85,7 @@ from lightrag.kg.shared_storage import (
     get_default_workspace,
     set_default_workspace,
     get_namespace_lock,
+    get_storage_keyed_lock,
 )
 
 from lightrag.base import (
@@ -1450,6 +1452,8 @@ class LightRAG(_RoleLLMMixin, _StorageMigrationMixin, _PipelineMixin):
         full_doc_id: str = None,
     ) -> None:
         update_storage = False
+        graph_lock_context = None
+        graph_lock_acquired = False
         try:
             # Insert chunks into vector storage
             all_chunks_data: dict[str, dict[str, str]] = {}
@@ -1489,6 +1493,
…[1709 chars omitted]…
 raise
         finally:
-            if update_storage:
-                await self._insert_done()
+            try:
+                if graph_lock_acquired and graph_lock_context is not None:
+                    await graph_lock_context.__aexit__(*sys.exc_info())
+            finally:
+                if update_storage:
+                    await self._insert_done()
 
     def query(
         self,

        Upsert an edge and its properties between two nodes identified by their labels.

        Args:
            source_node_id (str): Label of the source node (used as identifier)
            target_node_id (str): Label of the target node (used as identifier)
            edge_data (dict): dictionary of properties to set on the edge
        """
        # Keep all caller-controlled Cypher data in parameters. This includes
        # property names and values, which are encoded together as one agtype map;
        # an empty map is val
…[3530 chars omitted]…
eight`: "1.0"' in call["sql"]
    assert '`description`: "knows"' in call["sql"]


@pytest.mark.asyncio
async def test_upsert_edge_injection_payload():
    """Injection payloads in edge entity IDs are safely parameterized."""
    storage = make_graph_storage()
    injection_src = 'src"}) MATCH (x) DETACH DELETE x; //'
    injection_tgt = 'tgt"})-[r]-() DELETE r; //'
    captured_calls: list[dict] = []

        async def fake_query(sql, **kwargs):
            captured_sql.append(sql)
            return []
    
        with patch.object(storage, "_query", side_effect=fake_query):
            await storage.upsert_edge("Alice", "Bob", {"weight": "0.5"})
    
        sql = captured_sql[0
…[4276 chars omitted]…
========
FAILED tests/test_postgres_upsert_edge_cypher.py::test_upsert_edge_contains_optional_match_delete_create
FAILED tests/test_postgres_upsert_edge_cypher.py::test_upsert_edge_handles_empty_props
FAILED tests/test_postgres_cypher_injection.py::test_upsert_edge_uses_parameterized_cypher
FAILED tests/test_postgres_cypher_injection.py::test_upsert_edge_unicode_entity_ids
4 failed, 14 passed in 0.98s

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception(_is_transient_graph_write_error),
        reraise=True,
    )
    async def upsert_edge(
        self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
    ) -> None:
        """
        Upsert an edge and its properties between two nodes identified by their labels.

        Args:
            source_node_id (str): Label of the source node (use
…[1254 chars omitted]…
 $1::agtype) AS (r agtype)"
        )
        pg_params = {
            "params": json.dumps(
                {
                    "src_id": source_node_id,
                    "tgt_id": target_node_id,
                },
                ensure_ascii=False,
            ),
            "graph_name": self.graph_name,
        }
        # Advisory lock keys use the workspace graph and an order-independent

16-# ---------------------------------------------------------------------------
17-# Helpers
18-# ---------------------------------------------------------------------------
19-
20-
21-def make_graph_storage() -> PGGraphStorage:
22-    """Construct a PGGraphStorage instance with a mocked _query method."""
23-    storage = PGGraphStorage.__new__(PGGraphStorage)
24-    storage.workspace = "test_ws"
25-    storage.namespace = "test_graph"
26-    storage.graph_name = "test_graph"
27-    storage.db = MagicMock()
28-   
…[3537 chars omitted]…
   captured_calls.append({"sql": sql, **kwargs})
123-        return []
124-
--
187:# upsert_edge — parameterized Cypher
188-# ---------------------------------------------------------------------------
189-
190-
191-@pytest.mark.asyncio
192:async def test_upsert_edge_uses_parameterized_cypher():
193:    """upsert_edge must pass entity IDs as Cypher parameters."""
194-    storage = make_graph_storage()

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception(_is_transient_graph_write_error),
        reraise=True,
    )
    async def upsert_edge(
        self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
    ) -> None:
        """
        Upsert an edge and its properties between two nodes identified by their labels.

        Args:
            source_node_id (str): Label of the source node (use
…[7988 chars omitted]…
d in %.4fs",
                        timing_label,
                        time.perf_counter() - age_execute_start,
                    )

        except PGGraphQueryException:
            raise
        except Exception as e:
            if timing_label and not readonly:
                performance_timing_log(
                    "[%s] AGE execute failed after %.4fs",
                    timing_label,

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception(_is_transient_graph_write_error),
        reraise=True,
    )
    async def upsert_edge(
        self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
    ) -> None:
        """
        Upsert an edge and its properties between two nodes identified by their labels.

        Args:
            source_node_id (str): Label of the source node (use
…[986 chars omitted]…
ns independent of caller data.
        query = (
            f"SELECT * FROM cypher($2::text::name, "
            f"{_dollar_quote(cypher_query)}::cstring, $1::agtype) AS (r agtype)"
        )
        pg_params = {
            "params": json.dumps(
                {
                    "src_id": source_node_id,
                    "tgt_id": target_node_id,
                    "props": edge_data or {},

            "E'\\x01' || GREATEST($2::text, $3::text), 0))"
        )
        lock_params = (self.graph_name, source_node_id, target_node_id)
        timing_label = f"{self.workspace} PGGraphStorage.upsert_edge"
        total_start = time.perf_counter()
        performance_timing_log(
            "[%s] start source_node_id=%s target_node_id=%s",
            timing_label,
            source_node_id,

            return []
    
        with patch.object(storage, "_query", side_effect=fake_query):
            await storage.upsert_edge("Alice", "Bob", {"weight": "0.5"})
    
        sql = captured_sql[0]
        assert "OPTIONAL MATCH (source)-[old:DIRECTED]-(target)" in sql
        assert "DELETE ol
…[4177 chars omitted]…
========
FAILED tests/test_postgres_upsert_edge_cypher.py::test_upsert_edge_contains_optional_match_delete_create
FAILED tests/test_postgres_upsert_edge_cypher.py::test_upsert_edge_handles_empty_props
FAILED tests/test_postgres_cypher_injection.py::test_upsert_edge_uses_parameterized_cypher
FAILED tests/test_postgres_cypher_injection.py::test_upsert_edge_unicode_entity_ids
4 failed, 59 passed in 1.65s