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

15
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))
@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):
"""
@ -269,7 +271,7 @@ def create_query_routes(rag, api_key: Optional[str] = None, top_k: int = 60):
relationships = response.get("relationships", [])
chunks = response.get("chunks", [])
metadata = response.get("metadata", {})
# Validate data types
if not isinstance(entities, list):
entities = []
@ -279,12 +281,12 @@ def create_query_routes(rag, api_key: Optional[str] = None, top_k: int = 60):
chunks = []
if not isinstance(metadata, dict):
metadata = {}
return QueryDataResponse(
entities=entities,
relationships=relationships,
chunks=chunks,
metadata=metadata
metadata=metadata,
)
else:
# Fallback for unexpected response format
@ -292,7 +294,10 @@ def create_query_routes(rag, api_key: Optional[str] = None, top_k: int = 60):
entities=[],
relationships=[],
chunks=[],
metadata={"error": "Unexpected response format", "raw_response": str(response)}
metadata={
"error": "Unexpected response format",
"raw_response": str(response),
},
)
except Exception as e:
trace_exception(e)

4
lightrag/lightrag.py
View file

@ -2180,7 +2180,9 @@ class LightRAG:
entities_count = len(final_data.get("entities", []))
relationships_count = len(final_data.get("relationships", []))
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()
return final_data

4
lightrag/operate.py
View file

@ -2346,7 +2346,9 @@ async def kg_query(
if isinstance(context_result, tuple):
context, raw_data = context_result
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
else:
raise RuntimeError(

122
tests/test_aquery_data_endpoint.py
View file

@ -13,18 +13,15 @@ API_KEY = "your-secure-api-key-here-123"
BASE_URL = "http://localhost:9621"
# Unified authentication headers
AUTH_HEADERS = {
"Content-Type": "application/json",
"X-API-Key": API_KEY
}
AUTH_HEADERS = {"Content-Type": "application/json", "X-API-Key": API_KEY}
def test_aquery_data_endpoint():
"""Test the /query/data endpoint"""
# Use unified configuration
endpoint = f"{BASE_URL}/query/data"
# Query request
query_request = {
"query": "who authored LighRAG",
@ -34,12 +31,14 @@ def test_aquery_data_endpoint():
"max_entity_tokens": 4000,
"max_relation_tokens": 4000,
"max_total_tokens": 16000,
"enable_rerank": True
"enable_rerank": True,
}
print("=" * 60)
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("=" * 60)
print(f"Query content: {query_request['query']}")
@ -51,27 +50,24 @@ def test_aquery_data_endpoint():
# Send request
print("Sending request...")
start_time = time.time()
response = requests.post(
endpoint,
json=query_request,
headers=AUTH_HEADERS,
timeout=30
endpoint, json=query_request, headers=AUTH_HEADERS, timeout=30
)
end_time = time.time()
response_time = end_time - start_time
print(f"Response time: {response_time:.2f} seconds")
print(f"HTTP status code: {response.status_code}")
if response.status_code == 200:
data = response.json()
print_query_results(data)
else:
print(f"Request failed: {response.status_code}")
print(f"Error message: {response.text}")
except requests.exceptions.ConnectionError:
print("❌ Connection failed: Please ensure LightRAG API service is running")
print(" Start command: python -m lightrag.api.lightrag_server")
@ -83,91 +79,97 @@ def test_aquery_data_endpoint():
def print_query_results(data: Dict[str, Any]):
"""Format and print query results"""
entities = data.get("entities", [])
relationships = data.get("relationships", [])
chunks = data.get("chunks", [])
metadata = data.get("metadata", {})
print("\n📊 Query result statistics:")
print(f" Entity count: {len(entities)}")
print(f" Relationship count: {len(relationships)}")
print(f" Text chunk count: {len(chunks)}")
# Print metadata
if metadata:
print("\n🔍 Query metadata:")
print(f" Query mode: {metadata.get('query_mode', 'unknown')}")
keywords = metadata.get('keywords', {})
keywords = metadata.get("keywords", {})
if keywords:
high_level = keywords.get('high_level', [])
low_level = keywords.get('low_level', [])
high_level = keywords.get("high_level", [])
low_level = keywords.get("low_level", [])
if high_level:
print(f" High-level keywords: {', '.join(high_level)}")
if 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:
print(" Processing info:")
for key, value in processing_info.items():
print(f" {key}: {value}")
# Print entity information
if entities:
print("\n👥 Retrieved entities (first 5):")
for i, entity in enumerate(entities[:5]):
entity_name = entity.get('entity_name', 'Unknown')
entity_type = entity.get('entity_type', 'Unknown')
description = entity.get('description', 'No description')
file_path = entity.get('file_path', 'Unknown source')
entity_name = entity.get("entity_name", "Unknown")
entity_type = entity.get("entity_type", "Unknown")
description = entity.get("description", "No description")
file_path = entity.get("file_path", "Unknown source")
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()
# Print relationship information
if relationships:
print("🔗 Retrieved relationships (first 5):")
for i, rel in enumerate(relationships[:5]):
src = rel.get('src_id', 'Unknown')
tgt = rel.get('tgt_id', 'Unknown')
description = rel.get('description', 'No description')
keywords = rel.get('keywords', 'No keywords')
file_path = rel.get('file_path', 'Unknown source')
src = rel.get("src_id", "Unknown")
tgt = rel.get("tgt_id", "Unknown")
description = rel.get("description", "No description")
keywords = rel.get("keywords", "No keywords")
file_path = rel.get("file_path", "Unknown source")
print(f" {i+1}. {src} → {tgt}")
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()
# Print text chunk information
if chunks:
print("📄 Retrieved text chunks (first 3):")
for i, chunk in enumerate(chunks[:3]):
content = chunk.get('content', 'No content')
file_path = chunk.get('file_path', 'Unknown source')
chunk_id = chunk.get('chunk_id', 'Unknown ID')
content = chunk.get("content", "No content")
file_path = chunk.get("file_path", "Unknown source")
chunk_id = chunk.get("chunk_id", "Unknown ID")
print(f" {i+1}. Text chunk ID: {chunk_id}")
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("=" * 60)
def compare_with_regular_query():
"""Compare results between regular query and data query"""
query_text = "LightRAG的作者是谁"
print("\n🔄 Comparison test: Regular query vs Data query")
print("-" * 60)
# Regular query
try:
print("1. Regular query (/query):")
@ -175,13 +177,15 @@ def compare_with_regular_query():
f"{BASE_URL}/query",
json={"query": query_text, "mode": "mix"},
headers=AUTH_HEADERS,
timeout=30
timeout=30,
)
if regular_response.status_code == 200:
regular_data = regular_response.json()
response_text = regular_data.get('response', 'No response')
print(f" Generated answer: {response_text[:300]}{'...' if len(response_text) > 300 else ''}")
response_text = regular_data.get("response", "No response")
print(
f" Generated answer: {response_text[:300]}{'...' if len(response_text) > 300 else ''}"
)
else:
print(f" Regular query failed: {regular_response.status_code}")
if regular_response.status_code == 403:
@ -189,7 +193,7 @@ def compare_with_regular_query():
elif regular_response.status_code == 401:
print(" Unauthorized - Please check authentication information")
print(f" Error details: {regular_response.text}")
except Exception as e:
print(f" Regular query error: {str(e)}")
@ -197,10 +201,10 @@ def compare_with_regular_query():
if __name__ == "__main__":
# Run main test
test_aquery_data_endpoint()
# Run comparison test
compare_with_regular_query()
print("\n💡 Usage tips:")
print("1. Ensure LightRAG API service is running")
print("2. Adjust base_url and authentication information as needed")