hkuds-lightrag-3056-3129

config.ini.example
docker-build-push.sh
docker-compose-full.yml
docker-compose.podman.yml
docker-compose.yml
docs
env.docker-compose-full
env.example
examples
k8s-deploy
lightrag
lightrag.service.example
lightrag_hku.egg-info
lightrag_webui
prompts
pyproject.toml
reproduce
requirements-offline-llm.txt
requirements-offline-storage.txt
requirements-offline.txt
scripts
setup.py
tests
uv.lock
examples/opensearch_storage_demo.py
lightrag/base.py
lightrag/kg/memgraph_impl.py
lightrag/kg/mongo_impl.py
lightrag/kg/neo4j_impl.py
lightrag/kg/networkx_impl.py
lightrag/kg/opensearch_impl.py
lightrag/kg/postgres_impl.py
lightrag/kg/__init__.py
lightrag/lightrag.py

              JOIN b ON d.end_id   = b.vid
              LIMIT 1
            )
            OR EXISTS (
              SELECT 1
              FROM {self.graph_name}."DIRECTED" d
              JOIN a ON d.end_id   = a.vid
              JOIN b ON d.start_id = b.vid
              LIMIT 1
            ) AS edge_exists;
        """
        params = {
            "source_node_id": source_node_id,
            "target_node_id": target_node_id,
        }

…[8629 chars omitted]…
match the historical serial fallback.

        Args:
            edges: List of (source_node_id, target_node_id, edge_data) tuples.
        """
        if not edges:
            return
        deduped_edges: dict[tuple[str, str], tuple[str, str, dict[str, str]]] = {}
        for src, tgt, edge_data in edges:
            edge_key = tuple(sorted((src, tgt)))
            deduped_edges.pop(edge_key, None)

        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
        """
        # AGE does not support binding a full agtype map in ``SET r += $props``
        # (verified on AGE 1.5.0), and the inlined literal form ``SET r += {map}``
        # is also silently ignored for edges 
…[5320 chars omitted]…
f._query(query_nodes)
6451:                results = await self._query(query)
6527:        results = await self._query(query)
6569:        results = await self._query(query)
6619:            results = await self._query(query, params={"limit": limit})
6681:            results = await self._query(sql_query, params=dict(enumerate(params, 1)))
6704:            await self._query(drop_query, readonly=False)

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
…[3858 chars omitted]…
None,
                    multirows=True,
                    with_age=True,
                    graph_name=self.graph_name,
                    timing_label=timing_label,
                )
            else:
                age_execute_start = time.perf_counter()
                data = await self.db.execute(
                    query,
                    data=params,
                    upsert=upsert,

test_batch_embeddings.py
test_batch_graph_operations.py
test_bedrock_llm.py
test_chunk_options_persistence.py
test_chunker_recursive_character.py
test_chunker_semantic_vector.py
test_chunking.py
test_chunking_raw_lightrag_parity.py
test_content_hash_normalization.py
test_curl_aquery_data.sh
test_cwe89_opensearch_injection.py
test_degree_return_type.py
test_delete_file_variants_mineru_raw.py
test_description_api_validation.py
test_dimension_mismatch.py
test_doc_status_chunk_preservation.py
test_document_file_path_normalization.py
test_document_routes_docx_archive.py
test_document_routes_paginated.py
test_entity_extraction_stabi
…[471 chars omitted]…
test_milvus_index_config.py
test_milvus_index_creation.py
tests/test_batch_graph_operations.py
tests/test_chunk_options_persistence.py
tests/test_description_api_validation.py
tests/test_doc_status_chunk_preservation.py
tests/test_graph_storage.py
tests/test_opensearch_storage.py
tests/test_postgres_cypher_injection.py
tests/test_postgres_performance_timing.py
tests/test_postgres_upsert_edge_cypher.py

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"
   
…[3073 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"

                source_id = chunk_data["source_id"]
                file_path = normalize_document_file_path(
                    chunk_data.get("file_path", "custom_kg")
                )
                tokens = len(self.tokenizer.encode(chunk_content))
                chunk_order_index = (
                    0
                    if "chunk_order_index" not in chunk_data.keys()
                    else chunk_data["
…[8629 chars omitted]…
ystem_prompt: str | None = None,
    ) -> str | Iterator[str]:
        """
        Perform a sync query.

        Args:
            query (str): The query to be executed.
            param (QueryParam): Configuration parameters for query execution.
            prompt (Optional[str]): Custom prompts for fine-tuned control over the system's behavior. Defaults to None, which uses PROMPTS["rag_response"].

        entity_name: Name of the entity to delete
        entity_chunks_storage: Optional KV storage for tracking chunks that reference this entity
        relation_chunks_storage: Optional KV storage for tracking chunks that reference relations
    """
    # Use keyed lock for entity to ensure atomic graph and vector db operations
    workspace = entities_vdb.global_config.get("workspace", "")
    namespace = f"{workspace}
…[3175 chars omitted]…
          }

                    # Apply non-name updates first
                    if non_name_updates:
                        try:
                            logger.info(
                                "Entity Edit: applying non-name updates before merge"
                            )
                            await _edit_entity_impl(
                                chunk_entity_relation_graph,

    # ===== Phase 1: Process all entities concurrently =====
    log_message = f"Phase 1: Processing {total_entities_count} entities from {doc_id} (async: {graph_max_async})"
    logger.info(log_message)
    async with pipeline_status_lock:
        pipeline_status["latest_message"] = log_message
        pipeline_status["history_messages"].app
…[5910 chars omitted]…
                added_entities,  # Pass list to collect added entities
                        relation_chunks_storage,
                        entity_chunks_storage,  # Add entity_chunks_storage parameter
                    )

                    if edge_data is None:
                        return None, []

                    return edge_data, added_entities

                except Exception as e:

495:    async def _before_sleep(self, retry_state: RetryCallState) -> None:
506:    async def _run_with_retry(
572:    def _get_pool_snapshot(self) -> str:
613:    async def configure_vector_extension(self, connection: as
…[3285 chars omitted]…
delete_entity(self, entity_name: str) -> None:
3833:    async def delete_entity_relation(self, entity_name: str) -> None:
3855:    async def get_by_id(self, id: str) -> dict[str, Any] | None:
3878:    async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
3919:    async def get_vectors_by_ids(self, ids: list[str]) -> dict[str, list[float]]:
3971:    async def drop(self) -> dict[str, str]:

            if prepared_params:
                rows = await connection.fetch(sql, *prepared_params)
            else:
                rows = await connection.fetch(sql)
            fetch_elapsed = time.perf_counter() - fetch_start

            if timing_label:
                performance_timing_log(
                    "[%s] connection.fetch completed in %.4fs row_count=%s",
                    timing_label,
             
…[4893 chars omitted]…
t. 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,
    }
    _lock = asyncio.Lock()

    @staticmethod
    def get_config(vector_storage: str | None = None) -> dict[str, Any]:
        config = configparser.ConfigParser()

        self,
        operation: Callable[[asyncpg.Connection], Awaitable[T]],
        *,
        with_age: bool = False,
        graph_name: str | None = None,
        timing_label: str | None = None,
    ) -> T:
        """
        Execute a database operation with automatic retry for transient failures.

        Args:
            operation: Async callable that receives an active connection.
   
…[3655 chars omitted]…
aiting_count = len(getters)

        return (
            f"pool_state[min={min_size}, max={max_size}, holders={total_holders}, "
            f"acquired={acquired_count}, idle={idle_count}, waiting={waiting_count}]"
        )

    async def configure_vector_extension(self, connection: asyncpg.Connection) -> None:
        """Create VECTOR extension if it doesn't exist for vector similarity operations.

        except Exception as e:
            if timing_label and not readonly:
                performance_timing_log(
                    "[%s] AGE execute failed after %.4fs",
                    timing_label,
                    time.perf_counter() - age_execute_start,
                )
            raise PGGraphQueryException(
                {
                    "message": f"Error executing graph query: {query}",
                    "wrapped
…[79 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

794:        global_config["llm_cache_identities"] = {
800:        return global_config
805:        # `state` is None during the first _build_global_config() call from
818:        from lightrag.kg.shared_s
…[455 chars omitted]…
rage_cls, global_config=global_config
971:            global_config=global_config,
1040:            global_config=global_config,
1126:            from lightrag.kg.shared_storage import initialize_pipeline_status
1395:                global_config=self._build_global_config(),
1865:        global_config = self._build_global_config()
1899:                global_config,
1910:                global_config,

from lightrag.base import (
    BaseGraphStorage,
    BaseKVStorage,
    BaseVectorStorage,
    DocProcessingStatus,
    DocStatus,
    DocStatusStorage,
970:            workspace=self.workspace,
977:            workspace=self.workspace,
983:            workspace=self.workspace,
84:    get_default_workspace,
85:    set_default_workspace,
186:    workspace: str = field(default_factory=lambda: os.getenv("WORKSPACE", ""))
970:            workspace=self.workspace,
977:            workspace=self.workspace,
983:            workspace=self.workspace,
989:            workspace=self.workspace,
995:            workspace=self.workspace,
1001:            workspace=self.workspace,
1007:            workspace=self.workspace,
1013:            workspace=self.workspace,
1019:            workspace=self.workspace,
1025:            workspace=self.workspace,
1031:            workspace=self.workspace,
1039:            workspace=self.workspace,

102-    storage.graph_name = "test_graph"
103-    storage._query = AsyncMock(return_value=[])
104-
105:    await storage.upsert_edge(
106-        "node-1",
107-        "node-2",
108-        {
109-            "weight": 1.0,
110-            "description": "desc",
111-        },
112-    )
113-
114-    assert storage._query.await_args.kwargs["timing_label"] == (
115:        "test_ws PGGraphStorage.upsert_edge"
116-    )
117-
118-
119-def test_performance_timing_logs_reads_new_env_only(monkeypatch):
120-    with monkeypatch.context() as m
…[455 chars omitted]…
th passes params to db.execute
tests/test_postgres_cypher_injection.py:298:    """When readonly=False, _query must forward params to db.execute."""
tests/test_postgres_cypher_injection.py:306:    storage.db.execute = fake_execute
tests/test_postgres_halfvec.py:33:    db.execute = AsyncMock(side_effect=mock_execute)
tests/test_postgres_halfvec.py:96:        for call in mock_pg_db.execute.call_args_list

                f"[{self.workspace}] POSTGRES, upsert_node error on node_id: `{node_id}`"
            )
            raise

    @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 (used as identifier)
            target_node_id (str): Label of the target node (used as identifier)

                source_chunk_id = entity_data.get("source_id", "UNKNOWN")
                source_id = chunk_to_source_map.get(source_chunk_id, "UNKNOWN")
                file_path = normaliz
…[2021 chars omitted]…
_ids: set[str] = set()
            for relationship_data in deduped_relationships.values():
                needed_node_ids.add(relationship_data["src_id"])
                needed_node_ids.add(relationship_data["tgt_id"])

            existing_nodes = await self.chunk_entity_relation_graph.has_nodes_batch(
                list(needed_node_ids)
            )

            # Create missing nodes in batch

            )$
            if graph_lock_keys:$
                workspace = self.workspace or ""$
                graph_lock_namespace = f"{workspace}:GraphDB" if workspace else "GraphDB"$
                async with get_storage_keyed_lock($
                    graph_lock_keys,$
            
…[324 chars omitted]…
elationship storage is undirected, so keep only the last update$
                # for each endpoint pair regardless of order.$
                deduped_relationships: dict[tuple[str, str], dict[str, Any]] = {}$
                for relationship_data in custom_kg.get("relationships", []):$
                    src_id = relationship_data["src_id"]$
                    tgt_id = relationship_data["tgt_id"]$

                )
                async with get_storage_keyed_lock(
                    graph_lock_keys,
                    namespace=graph_lock_namespace,
                    enable_logging=False,
                ):
                    # Batch insert entities (reduces N serial awaits to 1)
                    if entity_nodes:
                        await self.chunk_entity_relation_graph.upsert_nodes_batch(entity_nodes)

  
…[2974 chars omitted]…
                                  "file_path": file_path,
                                            "created_at": int(time.time()),
                                        },
                                    )
                                )
                                existing_nodes.add(need_insert_id)

                        normalized_src_id, normalized_tgt_id = sorted((src_id, tgt_id))

transient recognized: True
Traceback (most recent call last):
  File "/repo/lightrag/kg/postgres_impl.py", line 5692, in upsert_edge
    await self.db.execute_transaction(
  File "/tmp/repro.py", line 45, in ve
    raise ValueError("boom")
ValueError: boom

The above except
…[2882 chars omitted]…
tity_id: $tgt_id})\n                     WITH source, target\n                     OPTIONAL MATCH (source)-[old:DIRECTED]-(target)\n                     DELETE old\n                     WITH source, target\n                     CREATE (source)-[r:DIRECTED {}]->(target)\n                     RETURN r$AGE1$::cstring, $1::agtype) AS (r agtype)', 'detail': "ValueError('boom')", 'error_type': 'ValueError'}

ERROR    lightrag:postgres_impl.py:5719 [test_ws] POSTGRES, upsert_edge error on edge: `NodeA`-`NodeB`
=========================== short test summary info ============================
FAILED tests/test_postgres_upsert_edge_cypher.py::test_upsert_edge_uses_delete_create_not_set
!!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 1 failures !!!!!!!!!!!!!!!!!!!!!!!!!!!
1 failed in 1.79s

10 failed, 33 passed, 6 skipped in 5.29s

    
            if isinstance(e.reason, ResponseError):
                raise RetryError(e, request=request)
    
            if isinstance(e.reason, _ProxyError):
                raise ProxyError(e, request=request)
    
            if isinstance(e.reason, _SSLError):
                # This branch is for urllib3 v1.22 and later.
                raise SSLError(e, request=request)
    
>           raise ConnectionEr
…[314 chars omitted]…
led to resolve 'openaipublic.blob.core.windows.net' ([Errno -3] Temporary failure 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
1 failed in 1.73s

no changes added to commit (use "git add" and/or "git commit -a")
Dropped refs/stash@{0} (4c94825c743e6b60ede2af0825de85755e23adfb)

        )

    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"' not in call["sql"].replace("$1::agtype", "")
    params = json.loads(call["params"]["params"])
    assert params["src_id"] == "Alice"
    assert params["tgt_id"] == "Bob"
    assert "props" not in params
    assert '`weight`: "1.0"' in call["sql"]
    assert '`description`: "knows"' in call["sql"]


@pytest.mark.asyncio

  File "/usr/local/lib/python3.12/asyncio/base_events.py", line 691, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/tmp/repro2.py", line 29, in main
    await rag.ainsert_custom_kg({
  File "/repo/lightrag/lightrag.py", line 1544, in ainsert_custom_kg
    for src_id, tgt_id in deduped_relationships
                          ^^^^^^^^^^^^^^^^^^^^^
UnboundLocalError: cannot access local variable 'deduped_relationships' where it is not associated with a value

  File "/usr/local/lib/python3.12/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^
…[138 chars omitted]…
         ^^^^^^^^^^^^^^^
  File "/tmp/repro2.py", line 43, in main
    await rag.ainsert_custom_kg({"chunks": [{"content": "x", "source_id": "s"}]})
  File "/repo/lightrag/lightrag.py", line 1677, in ainsert_custom_kg
    for dp in all_relationships_data
              ^^^^^^^^^^^^^^^^^^^^^^
UnboundLocalError: cannot access local variable 'all_relationships_data' where it is not associated with a value

ERROR tests/test_token_auto_renewal.py
!!!!!!!!!!!!!!!!!!! Interrupted: 12 errors during collection !!!!!!!!!!!!!!!!!!!
3 skipped, 1162 deselected, 12 errors in 36.82s

+        graph_name: str | None = None,
+        timing_label: str | None = None,
+    ) -> None:
+        """Execute multiple statements on the same connection inside one transaction.
+
+        All statements run against a single pooled connection wrapped in an
+        explicit transaction, so the group is atomic: either every statement
+        commits or none does. This is required for read-modify-write sequences
+        (e.g. lock + cypher upsert) 
…[5798 chars omitted]…
: f"Error executing graph edge upsert: {query}",
+                    "wrapped": query,
+                    "detail": repr(e),
+                    "error_type": e.__class__.__name__,
+                }
+            ) from e
 
     async def upsert_nodes_batch(self, nodes: list[tuple[str, dict[str, str]]]) -> None:
         """Batch insert/update multiple nodes while preserving input-order semantics.

@@ -1529,10 +1530,6 @@ class LightRAG(_RoleLLMMixin, _StorageMigrationMixin, _PipelineMixin):
                 all_entities_data.append(node_data_copy)
                 update_storage = True
 
-            # Batch insert entities (reduces N serial awaits to 1)
-            if entity_nodes:
-                await self.chunk_entity_relation_graph.upsert_nodes_batch(entity_nodes)
-
             # Relationship storage is undirected, so keep only the last update
             # for each endpoint pair regardless of order.
             deduped_relationshi
…[4860 chars omitted]…
                        )
-                        existing_nodes.add(need_insert_id)
 
-                normalized_src_id, normalized_tgt_id = sorted((src_id, tgt_id))
-
-                edge_data = {
-                    "weight": relationship_data.get("weight", 1.0),
-                    "description": relationship_data["description"],
-                    "keywords": relationship_data["keywords"],

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 | |     
…[8630 chars omitted]…
       merge_strategy: dict[str, str] = None,
3961 |         target_entity_data: dict[str, Any] = None,
     |                             ^^^^^^^^^^^^^^
3962 |     ) -> dict[str, Any]:
3963 |         loop = always_get_an_event_loop()
     |
help: Convert to `T | None`

Found 172 errors.
[*] 29 fixable with the `--fix` option (31 hidden fixes can be enabled with the `--unsafe-fixes` option).
ALL OK
OK