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,
|
1490 |
-
|
|
|
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": "
|
42 |
-
"port":
|
43 |
-
"model": "
|
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": "
|
767 |
-
"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": {
|