Roy commited on
Commit
1aa0610
·
1 Parent(s): ffa36eb

Enhance PostgreSQL vector storage with chunk_id support

Browse files

- 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 CHANGED
@@ -438,6 +438,7 @@ class PGVectorStorage(BaseVectorStorage):
438
  "entity_name": item["entity_name"],
439
  "content": item["content"],
440
  "content_vector": json.dumps(item["__vector__"].tolist()),
 
441
  }
442
  return upsert_sql, data
443
 
@@ -450,6 +451,7 @@ class PGVectorStorage(BaseVectorStorage):
450
  "target_id": item["tgt_id"],
451
  "content": item["content"],
452
  "content_vector": json.dumps(item["__vector__"].tolist()),
 
453
  }
454
  return upsert_sql, data
455
 
@@ -1486,8 +1488,9 @@ SQL_TEMPLATES = {
1486
  content_vector=EXCLUDED.content_vector,
1487
  update_time = CURRENT_TIMESTAMP
1488
  """,
1489
- "upsert_entity": """INSERT INTO LIGHTRAG_VDB_ENTITY (workspace, id, entity_name, content, content_vector)
1490
- VALUES ($1, $2, $3, $4, $5)
 
1491
  ON CONFLICT (workspace,id) DO UPDATE
1492
  SET entity_name=EXCLUDED.entity_name,
1493
  content=EXCLUDED.content,
@@ -1495,8 +1498,8 @@ SQL_TEMPLATES = {
1495
  update_time=CURRENT_TIMESTAMP
1496
  """,
1497
  "upsert_relationship": """INSERT INTO LIGHTRAG_VDB_RELATION (workspace, id, source_id,
1498
- target_id, content, content_vector)
1499
- VALUES ($1, $2, $3, $4, $5, $6)
1500
  ON CONFLICT (workspace,id) DO UPDATE
1501
  SET source_id=EXCLUDED.source_id,
1502
  target_id=EXCLUDED.target_id,
 
438
  "entity_name": item["entity_name"],
439
  "content": item["content"],
440
  "content_vector": json.dumps(item["__vector__"].tolist()),
441
+ "chunk_id": item["source_id"],
442
  }
443
  return upsert_sql, data
444
 
 
451
  "target_id": item["tgt_id"],
452
  "content": item["content"],
453
  "content_vector": json.dumps(item["__vector__"].tolist()),
454
+ "chunk_id": item["source_id"]
455
  }
456
  return upsert_sql, data
457
 
 
1488
  content_vector=EXCLUDED.content_vector,
1489
  update_time = CURRENT_TIMESTAMP
1490
  """,
1491
+ "upsert_entity": """INSERT INTO LIGHTRAG_VDB_ENTITY (workspace, id, entity_name, content,
1492
+ content_vector, chunk_id)
1493
+ VALUES ($1, $2, $3, $4, $5, $6)
1494
  ON CONFLICT (workspace,id) DO UPDATE
1495
  SET entity_name=EXCLUDED.entity_name,
1496
  content=EXCLUDED.content,
 
1498
  update_time=CURRENT_TIMESTAMP
1499
  """,
1500
  "upsert_relationship": """INSERT INTO LIGHTRAG_VDB_RELATION (workspace, id, source_id,
1501
+ target_id, content, content_vector, chunk_id)
1502
+ VALUES ($1, $2, $3, $4, $5, $6, $7)
1503
  ON CONFLICT (workspace,id) DO UPDATE
1504
  SET source_id=EXCLUDED.source_id,
1505
  target_id=EXCLUDED.target_id,
tests/test_lightrag_ollama_chat.py CHANGED
@@ -38,16 +38,16 @@ class McpError(Exception):
38
 
39
  DEFAULT_CONFIG = {
40
  "server": {
41
- "host": "localhost",
42
- "port": 9621,
43
- "model": "lightrag:latest",
44
  "timeout": 300,
45
  "max_retries": 1,
46
  "retry_delay": 1,
47
  },
48
  "test_cases": {
49
- "basic": {"query": "唐僧有几个徒弟"},
50
- "generate": {"query": "电视剧西游记导演是谁"},
51
  },
52
  }
53
 
@@ -763,8 +763,8 @@ def parse_args() -> argparse.Namespace:
763
  Configuration file (config.json):
764
  {
765
  "server": {
766
- "host": "localhost", # Server address
767
- "port": 9621, # Server port
768
  "model": "lightrag:latest" # Default model name
769
  },
770
  "test_cases": {
 
38
 
39
  DEFAULT_CONFIG = {
40
  "server": {
41
+ "host": "host.docker.internal",
42
+ "port": 11434,
43
+ "model": "llama3.2:latest",
44
  "timeout": 300,
45
  "max_retries": 1,
46
  "retry_delay": 1,
47
  },
48
  "test_cases": {
49
+ "basic": {"query": "How many disciples did Tang Seng have?"},
50
+ "generate": {"query": "Who directed the TV series Journey to the West?"},
51
  },
52
  }
53
 
 
763
  Configuration file (config.json):
764
  {
765
  "server": {
766
+ "host": "host.docker.internal", # Server address
767
+ "port": 11434, # Server port
768
  "model": "lightrag:latest" # Default model name
769
  },
770
  "test_cases": {