From bbe139cfebad995efc24a1a3e7f375530c6202cd Mon Sep 17 00:00:00 2001
From: Roy <arindamroy11235@gmail.com>
Date: Fri, 7 Mar 2025 20:18:01 +0000
Subject: [PATCH] Enhance PostgreSQL vector storage with chunk_id support

- Updated SQL templates for entity and relationship upsert to include chunk_id
- Modified PGVectorStorage methods to add chunk_id when inserting or updating records
- Expanded database schema to track chunk-level metadata
---
 lightrag/kg/postgres_impl.py       | 11 +++++++----
 tests/test_lightrag_ollama_chat.py | 14 +++++++-------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/lightrag/kg/postgres_impl.py b/lightrag/kg/postgres_impl.py
index a069cec0..3fc05f59 100644
--- a/lightrag/kg/postgres_impl.py
+++ b/lightrag/kg/postgres_impl.py
@@ -438,6 +438,7 @@ class PGVectorStorage(BaseVectorStorage):
             "entity_name": item["entity_name"],
             "content": item["content"],
             "content_vector": json.dumps(item["__vector__"].tolist()),
+            "chunk_id": item["source_id"],
         }
         return upsert_sql, data
 
@@ -450,6 +451,7 @@ class PGVectorStorage(BaseVectorStorage):
             "target_id": item["tgt_id"],
             "content": item["content"],
             "content_vector": json.dumps(item["__vector__"].tolist()),
+            "chunk_id": item["source_id"]
         }
         return upsert_sql, data
 
@@ -1486,8 +1488,9 @@ SQL_TEMPLATES = {
                       content_vector=EXCLUDED.content_vector,
                       update_time = CURRENT_TIMESTAMP
                      """,
-    "upsert_entity": """INSERT INTO LIGHTRAG_VDB_ENTITY (workspace, id, entity_name, content, content_vector)
-                      VALUES ($1, $2, $3, $4, $5)
+    "upsert_entity": """INSERT INTO LIGHTRAG_VDB_ENTITY (workspace, id, entity_name, content, 
+                      content_vector, chunk_id)
+                      VALUES ($1, $2, $3, $4, $5, $6)
                       ON CONFLICT (workspace,id) DO UPDATE
                       SET entity_name=EXCLUDED.entity_name,
                       content=EXCLUDED.content,
@@ -1495,8 +1498,8 @@ SQL_TEMPLATES = {
                       update_time=CURRENT_TIMESTAMP
                      """,
     "upsert_relationship": """INSERT INTO LIGHTRAG_VDB_RELATION (workspace, id, source_id,
-                      target_id, content, content_vector)
-                      VALUES ($1, $2, $3, $4, $5, $6)
+                      target_id, content, content_vector, chunk_id)
+                      VALUES ($1, $2, $3, $4, $5, $6, $7)
                       ON CONFLICT (workspace,id) DO UPDATE
                       SET source_id=EXCLUDED.source_id,
                       target_id=EXCLUDED.target_id,
diff --git a/tests/test_lightrag_ollama_chat.py b/tests/test_lightrag_ollama_chat.py
index 80038928..f19e2974 100644
--- a/tests/test_lightrag_ollama_chat.py
+++ b/tests/test_lightrag_ollama_chat.py
@@ -38,16 +38,16 @@ class McpError(Exception):
 
 DEFAULT_CONFIG = {
     "server": {
-        "host": "localhost",
-        "port": 9621,
-        "model": "lightrag:latest",
+        "host": "host.docker.internal",
+        "port": 11434,
+        "model": "llama3.2:latest",
         "timeout": 300,
         "max_retries": 1,
         "retry_delay": 1,
     },
     "test_cases": {
-        "basic": {"query": "唐僧有几个徒弟"},
-        "generate": {"query": "电视剧西游记导演是谁"},
+        "basic": {"query": "How many disciples did Tang Seng have?"},
+        "generate": {"query": "Who directed the TV series Journey to the West?"},
     },
 }
 
@@ -763,8 +763,8 @@ def parse_args() -> argparse.Namespace:
 Configuration file (config.json):
   {
     "server": {
-      "host": "localhost",      # Server address
-      "port": 9621,            # Server port
+      "host": "host.docker.internal",      # Server address
+      "port": 11434,            # Server port
       "model": "lightrag:latest" # Default model name
     },
     "test_cases": {