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

@@ -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,

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": {