gmakstutis/LightRAG
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix linting

Browse source
This commit is contained in:
yangdx 2025-09-15 02:59:21 +08:00
parent b1c8206346
commit c0d5abba6b
4 changed files with 79 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
lightrag/api/routers/query_routes.py
View file

@ -237,7 +237,9 @@ def create_query_routes(rag, api_key: Optional[str] = None, top_k: int = 60):
raise HTTPException(status_code=500, detail=str(e)) raise HTTPException(status_code=500, detail=str(e))
@router.post( @router.post(
"/query/data", response_model=QueryDataResponse, dependencies=[Depends(combined_auth)] "/query/data",
response_model=QueryDataResponse,
dependencies=[Depends(combined_auth)],
) )
async def query_data(request: QueryRequest): async def query_data(request: QueryRequest):
""" """
@ -284,7 +286,7 @@ def create_query_routes(rag, api_key: Optional[str] = None, top_k: int = 60):
entities=entities, entities=entities,
relationships=relationships, relationships=relationships,
chunks=chunks, chunks=chunks,
metadata=metadata metadata=metadata,
) )
else: else:
# Fallback for unexpected response format # Fallback for unexpected response format
@ -292,7 +294,10 @@ def create_query_routes(rag, api_key: Optional[str] = None, top_k: int = 60):
entities=[], entities=[],
relationships=[], relationships=[],
chunks=[], chunks=[],
metadata={"error": "Unexpected response format", "raw_response": str(response)} metadata={
"error": "Unexpected response format",
"raw_response": str(response),
},
) )
except Exception as e: except Exception as e:
trace_exception(e) trace_exception(e)

4
lightrag/lightrag.py
View file

@ -2180,7 +2180,9 @@ class LightRAG:
entities_count = len(final_data.get("entities", [])) entities_count = len(final_data.get("entities", []))
relationships_count = len(final_data.get("relationships", [])) relationships_count = len(final_data.get("relationships", []))
chunks_count = len(final_data.get("chunks", [])) chunks_count = len(final_data.get("chunks", []))
logger.debug(f"[aquery_data] Final result: {entities_count} entities, {relationships_count} relationships, {chunks_count} chunks") logger.debug(
f"[aquery_data] Final result: {entities_count} entities, {relationships_count} relationships, {chunks_count} chunks"
)
await self._query_done() await self._query_done()
return final_data return final_data

4
lightrag/operate.py
View file

@ -2346,7 +2346,9 @@ async def kg_query(
if isinstance(context_result, tuple): if isinstance(context_result, tuple):
context, raw_data = context_result context, raw_data = context_result
logger.debug(f"[kg_query] Context length: {len(context) if context else 0}") logger.debug(f"[kg_query] Context length: {len(context) if context else 0}")
logger.debug(f"[kg_query] Raw data entities: {len(raw_data.get('entities', []))}, relationships: {len(raw_data.get('relationships', []))}, chunks: {len(raw_data.get('chunks', []))}") logger.debug(
f"[kg_query] Raw data entities: {len(raw_data.get('entities', []))}, relationships: {len(raw_data.get('relationships', []))}, chunks: {len(raw_data.get('chunks', []))}"
)
return raw_data return raw_data
else: else:
raise RuntimeError( raise RuntimeError(

68
tests/test_aquery_data_endpoint.py
View file

@ -13,10 +13,7 @@ API_KEY = "your-secure-api-key-here-123"
BASE_URL = "http://localhost:9621" BASE_URL = "http://localhost:9621"
# Unified authentication headers # Unified authentication headers
AUTH_HEADERS = { AUTH_HEADERS = {"Content-Type": "application/json", "X-API-Key": API_KEY}
"Content-Type": "application/json",
"X-API-Key": API_KEY
}
def test_aquery_data_endpoint(): def test_aquery_data_endpoint():
@ -34,12 +31,14 @@ def test_aquery_data_endpoint():
"max_entity_tokens": 4000, "max_entity_tokens": 4000,
"max_relation_tokens": 4000, "max_relation_tokens": 4000,
"max_total_tokens": 16000, "max_total_tokens": 16000,
"enable_rerank": True "enable_rerank": True,
} }
print("=" * 60) print("=" * 60)
print("LightRAG aquery_data endpoint test") print("LightRAG aquery_data endpoint test")
print(" Returns structured data including entities, relationships and text chunks") print(
" Returns structured data including entities, relationships and text chunks"
)
print(" Can be used for custom processing and analysis") print(" Can be used for custom processing and analysis")
print("=" * 60) print("=" * 60)
print(f"Query content: {query_request['query']}") print(f"Query content: {query_request['query']}")
@ -53,10 +52,7 @@ def test_aquery_data_endpoint():
start_time = time.time() start_time = time.time()
response = requests.post( response = requests.post(
endpoint, endpoint, json=query_request, headers=AUTH_HEADERS, timeout=30
json=query_request,
headers=AUTH_HEADERS,
timeout=30
) )
end_time = time.time() end_time = time.time()
@ -99,16 +95,16 @@ def print_query_results(data: Dict[str, Any]):
print("\n🔍 Query metadata:") print("\n🔍 Query metadata:")
print(f" Query mode: {metadata.get('query_mode', 'unknown')}") print(f" Query mode: {metadata.get('query_mode', 'unknown')}")
keywords = metadata.get('keywords', {}) keywords = metadata.get("keywords", {})
if keywords: if keywords:
high_level = keywords.get('high_level', []) high_level = keywords.get("high_level", [])
low_level = keywords.get('low_level', []) low_level = keywords.get("low_level", [])
if high_level: if high_level:
print(f" High-level keywords: {', '.join(high_level)}") print(f" High-level keywords: {', '.join(high_level)}")
if low_level: if low_level:
print(f" Low-level keywords: {', '.join(low_level)}") print(f" Low-level keywords: {', '.join(low_level)}")
processing_info = metadata.get('processing_info', {}) processing_info = metadata.get("processing_info", {})
if processing_info: if processing_info:
print(" Processing info:") print(" Processing info:")
for key, value in processing_info.items(): for key, value in processing_info.items():
@ -118,13 +114,15 @@ def print_query_results(data: Dict[str, Any]):
if entities: if entities:
print("\n👥 Retrieved entities (first 5):") print("\n👥 Retrieved entities (first 5):")
for i, entity in enumerate(entities[:5]): for i, entity in enumerate(entities[:5]):
entity_name = entity.get('entity_name', 'Unknown') entity_name = entity.get("entity_name", "Unknown")
entity_type = entity.get('entity_type', 'Unknown') entity_type = entity.get("entity_type", "Unknown")
description = entity.get('description', 'No description') description = entity.get("description", "No description")
file_path = entity.get('file_path', 'Unknown source') file_path = entity.get("file_path", "Unknown source")
print(f" {i+1}. {entity_name} ({entity_type})") print(f" {i+1}. {entity_name} ({entity_type})")
print(f" Description: {description[:100]}{'...' if len(description) > 100 else ''}") print(
f" Description: {description[:100]}{'...' if len(description) > 100 else ''}"
)
print(f" Source: {file_path}") print(f" Source: {file_path}")
print() print()
@ -132,15 +130,17 @@ def print_query_results(data: Dict[str, Any]):
if relationships: if relationships:
print("🔗 Retrieved relationships (first 5):") print("🔗 Retrieved relationships (first 5):")
for i, rel in enumerate(relationships[:5]): for i, rel in enumerate(relationships[:5]):
src = rel.get('src_id', 'Unknown') src = rel.get("src_id", "Unknown")
tgt = rel.get('tgt_id', 'Unknown') tgt = rel.get("tgt_id", "Unknown")
description = rel.get('description', 'No description') description = rel.get("description", "No description")
keywords = rel.get('keywords', 'No keywords') keywords = rel.get("keywords", "No keywords")
file_path = rel.get('file_path', 'Unknown source') file_path = rel.get("file_path", "Unknown source")
print(f" {i+1}. {src} → {tgt}") print(f" {i+1}. {src} → {tgt}")
print(f" Keywords: {keywords}") print(f" Keywords: {keywords}")
print(f" Description: {description[:100]}{'...' if len(description) > 100 else ''}") print(
f" Description: {description[:100]}{'...' if len(description) > 100 else ''}"
)
print(f" Source: {file_path}") print(f" Source: {file_path}")
print() print()
@ -148,13 +148,15 @@ def print_query_results(data: Dict[str, Any]):
if chunks: if chunks:
print("📄 Retrieved text chunks (first 3):") print("📄 Retrieved text chunks (first 3):")
for i, chunk in enumerate(chunks[:3]): for i, chunk in enumerate(chunks[:3]):
content = chunk.get('content', 'No content') content = chunk.get("content", "No content")
file_path = chunk.get('file_path', 'Unknown source') file_path = chunk.get("file_path", "Unknown source")
chunk_id = chunk.get('chunk_id', 'Unknown ID') chunk_id = chunk.get("chunk_id", "Unknown ID")
print(f" {i+1}. Text chunk ID: {chunk_id}") print(f" {i+1}. Text chunk ID: {chunk_id}")
print(f" Source: {file_path}") print(f" Source: {file_path}")
print(f" Content: {content[:200]}{'...' if len(content) > 200 else ''}") print(
f" Content: {content[:200]}{'...' if len(content) > 200 else ''}"
)
print() print()
print("=" * 60) print("=" * 60)
@ -175,13 +177,15 @@ def compare_with_regular_query():
f"{BASE_URL}/query", f"{BASE_URL}/query",
json={"query": query_text, "mode": "mix"}, json={"query": query_text, "mode": "mix"},
headers=AUTH_HEADERS, headers=AUTH_HEADERS,
timeout=30 timeout=30,
) )
if regular_response.status_code == 200: if regular_response.status_code == 200:
regular_data = regular_response.json() regular_data = regular_response.json()
response_text = regular_data.get('response', 'No response') response_text = regular_data.get("response", "No response")
print(f" Generated answer: {response_text[:300]}{'...' if len(response_text) > 300 else ''}") print(
f" Generated answer: {response_text[:300]}{'...' if len(response_text) > 300 else ''}"
)
else: else:
print(f" Regular query failed: {regular_response.status_code}") print(f" Regular query failed: {regular_response.status_code}")
if regular_response.status_code == 403: if regular_response.status_code == 403: