gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update vector query methods to support ID filtering in PostgreSQL

Browse Source 
- Modified `mix_kg_vector_query` in operate.py to pass optional IDs to vector search
- Updated PostgreSQL SQL template to filter results using document IDs instead of chunk_id
- Improved query flexibility by allowing precise document selection during vector search
This commit is contained in:
Roy
2025-03-08 20:25:20 +00:00
parent 528fb11364
commit e31c0c8f6c
2 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
lightrag/kg/postgres_impl.py
View File

@@ -1597,7 +1597,7 @@ SQL_TEMPLATES = {
SELECT id, 1 - (content_vector <=> '[{embedding_string}]'::vector) as distance SELECT id, 1 - (content_vector <=> '[{embedding_string}]'::vector) as distance
FROM LIGHTRAG_DOC_CHUNKS FROM LIGHTRAG_DOC_CHUNKS
where workspace=$1 where workspace=$1
AND chunk_id IN (SELECT chunk_id FROM relevant_chunks) AND id IN (SELECT chunk_id FROM relevant_chunks)
) )
WHERE distance>$2 WHERE distance>$2
ORDER BY distance DESC ORDER BY distance DESC

3
lightrag/operate.py
View File

@@ -892,7 +892,8 @@ async def mix_kg_vector_query(
try: try:
# Reduce top_k for vector search in hybrid mode since we have structured information from KG # Reduce top_k for vector search in hybrid mode since we have structured information from KG
mix_topk = min(10, query_param.top_k) mix_topk = min(10, query_param.top_k)
results = await chunks_vdb.query(augmented_query, top_k=mix_topk) # TODO: add ids to the query
results = await chunks_vdb.query(augmented_query, top_k=mix_topk, ids = query_param.ids)
if not results: if not results:
return None return None