chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,781 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script: Demonstrates usage of aquery_data FastAPI endpoint
|
||||
Query content: Who is the author of LightRAG
|
||||
|
||||
Updated to handle the new data format where:
|
||||
- Response includes status, message, data, and metadata fields at top level
|
||||
- Actual query results (entities, relationships, chunks, references) are nested under 'data' field
|
||||
- Includes backward compatibility with legacy format
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
import time
|
||||
import json
|
||||
from typing import Dict, Any, List, Optional
|
||||
|
||||
# API configuration
|
||||
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}
|
||||
|
||||
|
||||
def validate_references_format(references: List[Dict[str, Any]]) -> bool:
|
||||
"""Validate the format of references list"""
|
||||
if not isinstance(references, list):
|
||||
print(f"❌ References should be a list, got {type(references)}")
|
||||
return False
|
||||
|
||||
for i, ref in enumerate(references):
|
||||
if not isinstance(ref, dict):
|
||||
print(f"❌ Reference {i} should be a dict, got {type(ref)}")
|
||||
return False
|
||||
|
||||
required_fields = ["reference_id", "file_path"]
|
||||
for field in required_fields:
|
||||
if field not in ref:
|
||||
print(f"❌ Reference {i} missing required field: {field}")
|
||||
return False
|
||||
|
||||
if not isinstance(ref[field], str):
|
||||
print(
|
||||
f"❌ Reference {i} field '{field}' should be string, got {type(ref[field])}"
|
||||
)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def parse_streaming_response(
|
||||
response_text: str,
|
||||
) -> tuple[Optional[List[Dict]], List[str], List[str]]:
|
||||
"""Parse streaming response and extract references, response chunks, and errors"""
|
||||
references = None
|
||||
response_chunks = []
|
||||
errors = []
|
||||
|
||||
lines = response_text.strip().split("\n")
|
||||
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if not line or line.startswith("data: "):
|
||||
if line.startswith("data: "):
|
||||
line = line[6:] # Remove 'data: ' prefix
|
||||
|
||||
if not line:
|
||||
continue
|
||||
|
||||
try:
|
||||
data = json.loads(line)
|
||||
|
||||
if "references" in data:
|
||||
references = data["references"]
|
||||
if "response" in data:
|
||||
response_chunks.append(data["response"])
|
||||
if "error" in data:
|
||||
errors.append(data["error"])
|
||||
|
||||
except json.JSONDecodeError:
|
||||
# Skip non-JSON lines (like SSE comments)
|
||||
continue
|
||||
|
||||
return references, response_chunks, errors
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.requires_api
|
||||
def test_query_endpoint_references():
|
||||
"""Test /query endpoint references functionality"""
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Testing /query endpoint references functionality")
|
||||
print("=" * 60)
|
||||
|
||||
query_text = "who authored LightRAG"
|
||||
endpoint = f"{BASE_URL}/query"
|
||||
|
||||
# Test 1: References enabled (default)
|
||||
print("\n🧪 Test 1: References enabled (default)")
|
||||
print("-" * 40)
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
endpoint,
|
||||
json={"query": query_text, "mode": "mix", "include_references": True},
|
||||
headers=AUTH_HEADERS,
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
|
||||
# Check response structure
|
||||
if "response" not in data:
|
||||
print("❌ Missing 'response' field")
|
||||
return False
|
||||
|
||||
if "references" not in data:
|
||||
print("❌ Missing 'references' field when include_references=True")
|
||||
return False
|
||||
|
||||
references = data["references"]
|
||||
if references is None:
|
||||
print("❌ References should not be None when include_references=True")
|
||||
return False
|
||||
|
||||
if not validate_references_format(references):
|
||||
return False
|
||||
|
||||
print(f"✅ References enabled: Found {len(references)} references")
|
||||
print(f" Response length: {len(data['response'])} characters")
|
||||
|
||||
# Display reference list
|
||||
if references:
|
||||
print(" 📚 Reference List:")
|
||||
for i, ref in enumerate(references, 1):
|
||||
ref_id = ref.get("reference_id", "Unknown")
|
||||
file_path = ref.get("file_path", "Unknown")
|
||||
print(f" {i}. ID: {ref_id} | File: {file_path}")
|
||||
|
||||
else:
|
||||
print(f"❌ Request failed: {response.status_code}")
|
||||
print(f" Error: {response.text}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Test 1 failed: {str(e)}")
|
||||
return False
|
||||
|
||||
# Test 2: References disabled
|
||||
print("\n🧪 Test 2: References disabled")
|
||||
print("-" * 40)
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
endpoint,
|
||||
json={"query": query_text, "mode": "mix", "include_references": False},
|
||||
headers=AUTH_HEADERS,
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
|
||||
# Check response structure
|
||||
if "response" not in data:
|
||||
print("❌ Missing 'response' field")
|
||||
return False
|
||||
|
||||
references = data.get("references")
|
||||
if references is not None:
|
||||
print("❌ References should be None when include_references=False")
|
||||
return False
|
||||
|
||||
print("✅ References disabled: No references field present")
|
||||
print(f" Response length: {len(data['response'])} characters")
|
||||
|
||||
else:
|
||||
print(f"❌ Request failed: {response.status_code}")
|
||||
print(f" Error: {response.text}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Test 2 failed: {str(e)}")
|
||||
return False
|
||||
|
||||
print("\n✅ /query endpoint references tests passed!")
|
||||
return True
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.requires_api
|
||||
def test_query_stream_endpoint_references():
|
||||
"""Test /query/stream endpoint references functionality"""
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Testing /query/stream endpoint references functionality")
|
||||
print("=" * 60)
|
||||
|
||||
query_text = "who authored LightRAG"
|
||||
endpoint = f"{BASE_URL}/query/stream"
|
||||
|
||||
# Test 1: Streaming with references enabled
|
||||
print("\n🧪 Test 1: Streaming with references enabled")
|
||||
print("-" * 40)
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
endpoint,
|
||||
json={"query": query_text, "mode": "mix", "include_references": True},
|
||||
headers=AUTH_HEADERS,
|
||||
timeout=30,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
# Collect streaming response
|
||||
full_response = ""
|
||||
for chunk in response.iter_content(chunk_size=1024, decode_unicode=True):
|
||||
if chunk:
|
||||
# Ensure chunk is string type
|
||||
if isinstance(chunk, bytes):
|
||||
chunk = chunk.decode("utf-8")
|
||||
full_response += chunk
|
||||
|
||||
# Parse streaming response
|
||||
references, response_chunks, errors = parse_streaming_response(
|
||||
full_response
|
||||
)
|
||||
|
||||
if errors:
|
||||
print(f"❌ Errors in streaming response: {errors}")
|
||||
return False
|
||||
|
||||
if references is None:
|
||||
print("❌ No references found in streaming response")
|
||||
return False
|
||||
|
||||
if not validate_references_format(references):
|
||||
return False
|
||||
|
||||
if not response_chunks:
|
||||
print("❌ No response chunks found in streaming response")
|
||||
return False
|
||||
|
||||
print(f"✅ Streaming with references: Found {len(references)} references")
|
||||
print(f" Response chunks: {len(response_chunks)}")
|
||||
print(
|
||||
f" Total response length: {sum(len(chunk) for chunk in response_chunks)} characters"
|
||||
)
|
||||
|
||||
# Display reference list
|
||||
if references:
|
||||
print(" 📚 Reference List:")
|
||||
for i, ref in enumerate(references, 1):
|
||||
ref_id = ref.get("reference_id", "Unknown")
|
||||
file_path = ref.get("file_path", "Unknown")
|
||||
print(f" {i}. ID: {ref_id} | File: {file_path}")
|
||||
|
||||
else:
|
||||
print(f"❌ Request failed: {response.status_code}")
|
||||
print(f" Error: {response.text}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Test 1 failed: {str(e)}")
|
||||
return False
|
||||
|
||||
# Test 2: Streaming with references disabled
|
||||
print("\n🧪 Test 2: Streaming with references disabled")
|
||||
print("-" * 40)
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
endpoint,
|
||||
json={"query": query_text, "mode": "mix", "include_references": False},
|
||||
headers=AUTH_HEADERS,
|
||||
timeout=30,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
# Collect streaming response
|
||||
full_response = ""
|
||||
for chunk in response.iter_content(chunk_size=1024, decode_unicode=True):
|
||||
if chunk:
|
||||
# Ensure chunk is string type
|
||||
if isinstance(chunk, bytes):
|
||||
chunk = chunk.decode("utf-8")
|
||||
full_response += chunk
|
||||
|
||||
# Parse streaming response
|
||||
references, response_chunks, errors = parse_streaming_response(
|
||||
full_response
|
||||
)
|
||||
|
||||
if errors:
|
||||
print(f"❌ Errors in streaming response: {errors}")
|
||||
return False
|
||||
|
||||
if references is not None:
|
||||
print("❌ References should be None when include_references=False")
|
||||
return False
|
||||
|
||||
if not response_chunks:
|
||||
print("❌ No response chunks found in streaming response")
|
||||
return False
|
||||
|
||||
print("✅ Streaming without references: No references present")
|
||||
print(f" Response chunks: {len(response_chunks)}")
|
||||
print(
|
||||
f" Total response length: {sum(len(chunk) for chunk in response_chunks)} characters"
|
||||
)
|
||||
|
||||
else:
|
||||
print(f"❌ Request failed: {response.status_code}")
|
||||
print(f" Error: {response.text}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Test 2 failed: {str(e)}")
|
||||
return False
|
||||
|
||||
print("\n✅ /query/stream endpoint references tests passed!")
|
||||
return True
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.requires_api
|
||||
def test_references_consistency():
|
||||
"""Test references consistency across all endpoints"""
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Testing references consistency across endpoints")
|
||||
print("=" * 60)
|
||||
|
||||
query_text = "who authored LightRAG"
|
||||
query_params = {
|
||||
"query": query_text,
|
||||
"mode": "mix",
|
||||
"top_k": 10,
|
||||
"chunk_top_k": 8,
|
||||
"include_references": True,
|
||||
}
|
||||
|
||||
references_data = {}
|
||||
|
||||
# Test /query endpoint
|
||||
print("\n🧪 Testing /query endpoint")
|
||||
print("-" * 40)
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{BASE_URL}/query", json=query_params, headers=AUTH_HEADERS, timeout=30
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
references_data["query"] = data.get("references", [])
|
||||
print(f"✅ /query: {len(references_data['query'])} references")
|
||||
else:
|
||||
print(f"❌ /query failed: {response.status_code}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ /query test failed: {str(e)}")
|
||||
return False
|
||||
|
||||
# Test /query/stream endpoint
|
||||
print("\n🧪 Testing /query/stream endpoint")
|
||||
print("-" * 40)
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{BASE_URL}/query/stream",
|
||||
json=query_params,
|
||||
headers=AUTH_HEADERS,
|
||||
timeout=30,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
full_response = ""
|
||||
for chunk in response.iter_content(chunk_size=1024, decode_unicode=True):
|
||||
if chunk:
|
||||
# Ensure chunk is string type
|
||||
if isinstance(chunk, bytes):
|
||||
chunk = chunk.decode("utf-8")
|
||||
full_response += chunk
|
||||
|
||||
references, _, errors = parse_streaming_response(full_response)
|
||||
|
||||
if errors:
|
||||
print(f"❌ Errors: {errors}")
|
||||
return False
|
||||
|
||||
references_data["stream"] = references or []
|
||||
print(f"✅ /query/stream: {len(references_data['stream'])} references")
|
||||
else:
|
||||
print(f"❌ /query/stream failed: {response.status_code}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ /query/stream test failed: {str(e)}")
|
||||
return False
|
||||
|
||||
# Test /query/data endpoint
|
||||
print("\n🧪 Testing /query/data endpoint")
|
||||
print("-" * 40)
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{BASE_URL}/query/data",
|
||||
json=query_params,
|
||||
headers=AUTH_HEADERS,
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
query_data = data.get("data", {})
|
||||
references_data["data"] = query_data.get("references", [])
|
||||
print(f"✅ /query/data: {len(references_data['data'])} references")
|
||||
else:
|
||||
print(f"❌ /query/data failed: {response.status_code}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ /query/data test failed: {str(e)}")
|
||||
return False
|
||||
|
||||
# Compare references consistency
|
||||
print("\n🔍 Comparing references consistency")
|
||||
print("-" * 40)
|
||||
|
||||
# Convert to sets of (reference_id, file_path) tuples for comparison
|
||||
def refs_to_set(refs):
|
||||
return set(
|
||||
(ref.get("reference_id", ""), ref.get("file_path", "")) for ref in refs
|
||||
)
|
||||
|
||||
query_refs = refs_to_set(references_data["query"])
|
||||
stream_refs = refs_to_set(references_data["stream"])
|
||||
data_refs = refs_to_set(references_data["data"])
|
||||
|
||||
# Check consistency
|
||||
consistency_passed = True
|
||||
|
||||
if query_refs != stream_refs:
|
||||
print("❌ References mismatch between /query and /query/stream")
|
||||
print(f" /query only: {query_refs - stream_refs}")
|
||||
print(f" /query/stream only: {stream_refs - query_refs}")
|
||||
consistency_passed = False
|
||||
|
||||
if query_refs != data_refs:
|
||||
print("❌ References mismatch between /query and /query/data")
|
||||
print(f" /query only: {query_refs - data_refs}")
|
||||
print(f" /query/data only: {data_refs - query_refs}")
|
||||
consistency_passed = False
|
||||
|
||||
if stream_refs != data_refs:
|
||||
print("❌ References mismatch between /query/stream and /query/data")
|
||||
print(f" /query/stream only: {stream_refs - data_refs}")
|
||||
print(f" /query/data only: {data_refs - stream_refs}")
|
||||
consistency_passed = False
|
||||
|
||||
if consistency_passed:
|
||||
print("✅ All endpoints return consistent references")
|
||||
print(f" Common references count: {len(query_refs)}")
|
||||
|
||||
# Display common reference list
|
||||
if query_refs:
|
||||
print(" 📚 Common Reference List:")
|
||||
for i, (ref_id, file_path) in enumerate(sorted(query_refs), 1):
|
||||
print(f" {i}. ID: {ref_id} | File: {file_path}")
|
||||
|
||||
return consistency_passed
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.requires_api
|
||||
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",
|
||||
"mode": "mix", # Use mixed mode to get the most comprehensive results
|
||||
"top_k": 20,
|
||||
"chunk_top_k": 15,
|
||||
"max_entity_tokens": 4000,
|
||||
"max_relation_tokens": 4000,
|
||||
"max_total_tokens": 16000,
|
||||
"enable_rerank": True,
|
||||
}
|
||||
|
||||
print("=" * 60)
|
||||
print("LightRAG aquery_data endpoint test")
|
||||
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']}")
|
||||
print(f"Query mode: {query_request['mode']}")
|
||||
print(f"API endpoint: {endpoint}")
|
||||
print("-" * 60)
|
||||
|
||||
try:
|
||||
# Send request
|
||||
print("Sending request...")
|
||||
start_time = time.time()
|
||||
|
||||
response = requests.post(
|
||||
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")
|
||||
except requests.exceptions.Timeout:
|
||||
print("❌ Request timeout: Query processing took too long")
|
||||
except Exception as e:
|
||||
print(f"❌ Error occurred: {str(e)}")
|
||||
|
||||
|
||||
def print_query_results(data: Dict[str, Any]):
|
||||
"""Format and print query results"""
|
||||
|
||||
# Check for new data format with status and message
|
||||
status = data.get("status", "unknown")
|
||||
message = data.get("message", "")
|
||||
|
||||
print(f"\n📋 Query Status: {status}")
|
||||
if message:
|
||||
print(f"📋 Message: {message}")
|
||||
|
||||
# Handle new nested data format
|
||||
query_data = data.get("data", {})
|
||||
|
||||
# Fallback to old format if new format is not present
|
||||
if not query_data and any(
|
||||
key in data for key in ["entities", "relationships", "chunks"]
|
||||
):
|
||||
print(" (Using legacy data format)")
|
||||
query_data = data
|
||||
|
||||
entities = query_data.get("entities", [])
|
||||
relationships = query_data.get("relationships", [])
|
||||
chunks = query_data.get("chunks", [])
|
||||
references = query_data.get("references", [])
|
||||
|
||||
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(f" Reference count: {len(references)}")
|
||||
|
||||
# Print metadata (now at top level in new format)
|
||||
metadata = data.get("metadata", {})
|
||||
if metadata:
|
||||
print("\n🔍 Query metadata:")
|
||||
print(f" Query mode: {metadata.get('query_mode', 'unknown')}")
|
||||
|
||||
keywords = metadata.get("keywords", {})
|
||||
if keywords:
|
||||
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", {})
|
||||
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")
|
||||
reference_id = entity.get("reference_id", "No reference")
|
||||
|
||||
print(f" {i + 1}. {entity_name} ({entity_type})")
|
||||
print(
|
||||
f" Description: {description[:100]}{'...' if len(description) > 100 else ''}"
|
||||
)
|
||||
print(f" Source: {file_path}")
|
||||
print(f" Reference ID: {reference_id}")
|
||||
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")
|
||||
reference_id = rel.get("reference_id", "No reference")
|
||||
|
||||
print(f" {i + 1}. {src} → {tgt}")
|
||||
print(f" Keywords: {keywords}")
|
||||
print(
|
||||
f" Description: {description[:100]}{'...' if len(description) > 100 else ''}"
|
||||
)
|
||||
print(f" Source: {file_path}")
|
||||
print(f" Reference ID: {reference_id}")
|
||||
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")
|
||||
reference_id = chunk.get("reference_id", "No reference")
|
||||
|
||||
print(f" {i + 1}. Text chunk ID: {chunk_id}")
|
||||
print(f" Source: {file_path}")
|
||||
print(f" Reference ID: {reference_id}")
|
||||
print(
|
||||
f" Content: {content[:200]}{'...' if len(content) > 200 else ''}"
|
||||
)
|
||||
print()
|
||||
|
||||
# Print references information (new in updated format)
|
||||
if references:
|
||||
print("📚 References:")
|
||||
for i, ref in enumerate(references):
|
||||
reference_id = ref.get("reference_id", "Unknown ID")
|
||||
file_path = ref.get("file_path", "Unknown source")
|
||||
print(f" {i + 1}. Reference ID: {reference_id}")
|
||||
print(f" File Path: {file_path}")
|
||||
print()
|
||||
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.requires_api
|
||||
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):")
|
||||
regular_response = requests.post(
|
||||
f"{BASE_URL}/query",
|
||||
json={"query": query_text, "mode": "mix"},
|
||||
headers=AUTH_HEADERS,
|
||||
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 ''}"
|
||||
)
|
||||
else:
|
||||
print(f" Regular query failed: {regular_response.status_code}")
|
||||
if regular_response.status_code == 403:
|
||||
print(" Authentication failed - Please check API Key configuration")
|
||||
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)}")
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.requires_api
|
||||
def run_all_reference_tests():
|
||||
"""Run all reference-related tests"""
|
||||
|
||||
print("\n" + "🚀" * 20)
|
||||
print("LightRAG References Test Suite")
|
||||
print("🚀" * 20)
|
||||
|
||||
all_tests_passed = True
|
||||
|
||||
# Test 1: /query endpoint references
|
||||
try:
|
||||
if not test_query_endpoint_references():
|
||||
all_tests_passed = False
|
||||
except Exception as e:
|
||||
print(f"❌ /query endpoint test failed with exception: {str(e)}")
|
||||
all_tests_passed = False
|
||||
|
||||
# Test 2: /query/stream endpoint references
|
||||
try:
|
||||
if not test_query_stream_endpoint_references():
|
||||
all_tests_passed = False
|
||||
except Exception as e:
|
||||
print(f"❌ /query/stream endpoint test failed with exception: {str(e)}")
|
||||
all_tests_passed = False
|
||||
|
||||
# Test 3: References consistency across endpoints
|
||||
try:
|
||||
if not test_references_consistency():
|
||||
all_tests_passed = False
|
||||
except Exception as e:
|
||||
print(f"❌ References consistency test failed with exception: {str(e)}")
|
||||
all_tests_passed = False
|
||||
|
||||
# Final summary
|
||||
print("\n" + "=" * 60)
|
||||
print("TEST SUITE SUMMARY")
|
||||
print("=" * 60)
|
||||
|
||||
if all_tests_passed:
|
||||
print("🎉 ALL TESTS PASSED!")
|
||||
print("✅ /query endpoint references functionality works correctly")
|
||||
print("✅ /query/stream endpoint references functionality works correctly")
|
||||
print("✅ References are consistent across all endpoints")
|
||||
print("\n🔧 System is ready for production use with reference support!")
|
||||
else:
|
||||
print("❌ SOME TESTS FAILED!")
|
||||
print("Please check the error messages above and fix the issues.")
|
||||
print("\n🔧 System needs attention before production deployment.")
|
||||
|
||||
return all_tests_passed
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
if len(sys.argv) > 1 and sys.argv[1] == "--references-only":
|
||||
# Run only the new reference tests
|
||||
success = run_all_reference_tests()
|
||||
sys.exit(0 if success else 1)
|
||||
else:
|
||||
# Run original tests plus new reference tests
|
||||
print("Running original aquery_data endpoint test...")
|
||||
test_aquery_data_endpoint()
|
||||
|
||||
print("\nRunning comparison test...")
|
||||
compare_with_regular_query()
|
||||
|
||||
print("\nRunning new reference tests...")
|
||||
run_all_reference_tests()
|
||||
|
||||
print("\n💡 Usage tips:")
|
||||
print("1. Ensure LightRAG API service is running")
|
||||
print("2. Adjust base_url and authentication information as needed")
|
||||
print("3. Modify query parameters to test different retrieval strategies")
|
||||
print("4. Data query results can be used for further analysis and processing")
|
||||
print("5. Run with --references-only flag to test only reference functionality")
|
||||
Executable
+271
@@ -0,0 +1,271 @@
|
||||
#!/bin/bash
|
||||
|
||||
# LightRAG aquery_data endpoint test script
|
||||
# Use curl command to test the new /query/data endpoint and validate the new data format
|
||||
|
||||
echo "🚀 LightRAG aquery_data Endpoint Test (New Data Format Validation)"
|
||||
echo "=================================================="
|
||||
|
||||
# Base URL (adjust according to actual deployment)
|
||||
BASE_URL="http://localhost:9621"
|
||||
|
||||
# Color definitions
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
YELLOW='\033[1;33m'
|
||||
BLUE='\033[0;34m'
|
||||
NC='\033[0m' # No Color
|
||||
|
||||
# Test result statistics
|
||||
TOTAL_TESTS=0
|
||||
PASSED_TESTS=0
|
||||
FAILED_TESTS=0
|
||||
|
||||
# Function to validate success response format
|
||||
validate_success_response() {
|
||||
local response="$1"
|
||||
local test_name="$2"
|
||||
local expected_mode="$3"
|
||||
|
||||
echo -e "${BLUE}Validating $test_name response format...${NC}"
|
||||
|
||||
# Check if valid JSON
|
||||
if ! echo "$response" | jq . >/dev/null 2>&1; then
|
||||
echo -e "${RED}❌ Response is not valid JSON format${NC}"
|
||||
return 1
|
||||
fi
|
||||
|
||||
# Validate required fields
|
||||
local status=$(echo "$response" | jq -r '.status // "missing"')
|
||||
local message=$(echo "$response" | jq -r '.message // "missing"')
|
||||
local data_exists=$(echo "$response" | jq 'has("data")')
|
||||
local metadata_exists=$(echo "$response" | jq 'has("metadata")')
|
||||
|
||||
echo " Status: $status"
|
||||
echo " Message: $message"
|
||||
|
||||
# Validate data structure
|
||||
if [[ "$data_exists" == "true" ]]; then
|
||||
local entities_count=$(echo "$response" | jq '.data.entities | length // 0')
|
||||
local relationships_count=$(echo "$response" | jq '.data.relationships | length // 0')
|
||||
local chunks_count=$(echo "$response" | jq '.data.chunks | length // 0')
|
||||
local references_count=$(echo "$response" | jq '.data.references | length // 0')
|
||||
|
||||
echo " Data.entities: $entities_count"
|
||||
echo " Data.relationships: $relationships_count"
|
||||
echo " Data.chunks: $chunks_count"
|
||||
echo " Data.references: $references_count"
|
||||
else
|
||||
echo -e "${RED} ❌ Missing 'data' field${NC}"
|
||||
return 1
|
||||
fi
|
||||
|
||||
# Validate metadata
|
||||
if [[ "$metadata_exists" == "true" ]]; then
|
||||
local query_mode=$(echo "$response" | jq -r '.metadata.query_mode // "missing"')
|
||||
local keywords_exists=$(echo "$response" | jq 'has("metadata") and (.metadata | has("keywords"))')
|
||||
local processing_info_exists=$(echo "$response" | jq 'has("metadata") and (.metadata | has("processing_info"))')
|
||||
|
||||
echo " Metadata.query_mode: $query_mode"
|
||||
echo " Metadata.keywords: $keywords_exists"
|
||||
echo " Metadata.processing_info: $processing_info_exists"
|
||||
|
||||
# Validate if query mode matches
|
||||
if [[ "$expected_mode" != "" && "$query_mode" != "$expected_mode" ]]; then
|
||||
echo -e "${YELLOW} ⚠️ Query mode mismatch: expected '$expected_mode', actual '$query_mode'${NC}"
|
||||
fi
|
||||
else
|
||||
echo -e "${RED} ❌ Missing 'metadata' field${NC}"
|
||||
return 1
|
||||
fi
|
||||
|
||||
# Validate status
|
||||
if [[ "$status" == "success" ]]; then
|
||||
echo -e "${GREEN} ✅ Response format validation passed${NC}"
|
||||
return 0
|
||||
else
|
||||
echo -e "${RED} ❌ Status is not 'success': $status${NC}"
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# Function to validate error response format
|
||||
validate_error_response() {
|
||||
local response="$1"
|
||||
local test_name="$2"
|
||||
|
||||
echo -e "${BLUE}Validating $test_name response format...${NC}"
|
||||
|
||||
# Check if valid JSON
|
||||
if ! echo "$response" | jq . >/dev/null 2>&1; then
|
||||
echo -e "${RED}❌ Response is not valid JSON format${NC}"
|
||||
return 1
|
||||
fi
|
||||
|
||||
# Validate required fields
|
||||
local status=$(echo "$response" | jq -r '.status // "missing"')
|
||||
local message=$(echo "$response" | jq -r '.message // "missing"')
|
||||
local data_exists=$(echo "$response" | jq 'has("data")')
|
||||
local metadata_exists=$(echo "$response" | jq 'has("metadata")')
|
||||
|
||||
echo " Status: $status"
|
||||
echo " Message: $message"
|
||||
|
||||
# Validate basic structure exists
|
||||
if [[ "$data_exists" != "true" ]]; then
|
||||
echo -e "${RED} ❌ Missing 'data' field${NC}"
|
||||
return 1
|
||||
fi
|
||||
|
||||
if [[ "$metadata_exists" != "true" ]]; then
|
||||
echo -e "${RED} ❌ Missing 'metadata' field${NC}"
|
||||
return 1
|
||||
fi
|
||||
|
||||
echo " Data: {}"
|
||||
echo " Metadata: {}"
|
||||
|
||||
# Validate status should be failure
|
||||
if [[ "$status" == "failure" ]]; then
|
||||
echo -e "${GREEN} ✅ Error response format validation passed${NC}"
|
||||
return 0
|
||||
else
|
||||
echo -e "${RED} ❌ Status is not 'failure': $status${NC}"
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# Function to run success test
|
||||
run_success_test() {
|
||||
local test_name="$1"
|
||||
local query_data="$2"
|
||||
local expected_mode="$3"
|
||||
local print_json="${4:-false}" # Optional parameter: whether to print JSON response (default: false)
|
||||
|
||||
echo ""
|
||||
echo "=================================="
|
||||
echo -e "${BLUE}$test_name${NC}"
|
||||
echo "=================================="
|
||||
|
||||
TOTAL_TESTS=$((TOTAL_TESTS + 1))
|
||||
|
||||
# Send request
|
||||
echo "Sending request..."
|
||||
local response=$(curl -s -X POST "${BASE_URL}/query/data" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "X-API-Key: your-secure-api-key-here-123" \
|
||||
-d "$query_data")
|
||||
|
||||
# Check if curl succeeded
|
||||
if [[ $? -ne 0 ]]; then
|
||||
echo -e "${RED}❌ Request failed - cannot connect to server${NC}"
|
||||
FAILED_TESTS=$((FAILED_TESTS + 1))
|
||||
return 1
|
||||
fi
|
||||
|
||||
# Print JSON response if requested
|
||||
if [[ "$print_json" == "true" ]]; then
|
||||
echo ""
|
||||
echo "Response JSON:"
|
||||
echo "$response" | jq '.' 2>/dev/null || echo "$response"
|
||||
echo ""
|
||||
fi
|
||||
|
||||
# Validate response
|
||||
if validate_success_response "$response" "$test_name" "$expected_mode"; then
|
||||
PASSED_TESTS=$((PASSED_TESTS + 1))
|
||||
echo -e "${GREEN}✅ $test_name test passed${NC}"
|
||||
else
|
||||
FAILED_TESTS=$((FAILED_TESTS + 1))
|
||||
echo -e "${RED}❌ $test_name test failed${NC}"
|
||||
echo "Raw response:"
|
||||
echo "$response" | jq '.' 2>/dev/null || echo "$response"
|
||||
fi
|
||||
}
|
||||
|
||||
# Function to run error test
|
||||
run_error_test() {
|
||||
local test_name="$1"
|
||||
local query_data="$2"
|
||||
|
||||
echo ""
|
||||
echo "=================================="
|
||||
echo -e "${BLUE}$test_name${NC}"
|
||||
echo "=================================="
|
||||
|
||||
TOTAL_TESTS=$((TOTAL_TESTS + 1))
|
||||
|
||||
# Send request
|
||||
echo "Sending request..."
|
||||
local response=$(curl -s -X POST "${BASE_URL}/query/data" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "X-API-Key: your-secure-api-key-here-123" \
|
||||
-d "$query_data")
|
||||
|
||||
# Check if curl succeeded
|
||||
if [[ $? -ne 0 ]]; then
|
||||
echo -e "${RED}❌ Request failed - cannot connect to server${NC}"
|
||||
FAILED_TESTS=$((FAILED_TESTS + 1))
|
||||
return 1
|
||||
fi
|
||||
|
||||
# Validate response
|
||||
if validate_error_response "$response" "$test_name"; then
|
||||
PASSED_TESTS=$((PASSED_TESTS + 1))
|
||||
echo -e "${GREEN}✅ $test_name test passed${NC}"
|
||||
else
|
||||
FAILED_TESTS=$((FAILED_TESTS + 1))
|
||||
echo -e "${RED}❌ $test_name test failed${NC}"
|
||||
echo "Raw response:"
|
||||
echo "$response" | jq '.' 2>/dev/null || echo "$response"
|
||||
fi
|
||||
}
|
||||
|
||||
# Start tests
|
||||
echo "Starting tests for new /query/data endpoint data format..."
|
||||
echo ""
|
||||
|
||||
# Test 1: Basic query test (mix mode)
|
||||
run_success_test "1. Basic Query Test (mix mode)" '{
|
||||
"query": "What is GraphRAG",
|
||||
"mode": "mix",
|
||||
"top_k": 5
|
||||
}' "mix" "true" # Output full JSON
|
||||
|
||||
# Test 2: Detailed parameter query test (hybrid mode)
|
||||
run_success_test "2. Detailed Parameter Query Test (hybrid mode)" '{
|
||||
"query": "What is GraphRAG",
|
||||
"mode": "hybrid",
|
||||
"top_k": 5,
|
||||
"chunk_top_k": 8,
|
||||
"max_entity_tokens": 4000,
|
||||
"max_relation_tokens": 4000,
|
||||
"max_total_tokens": 16000,
|
||||
"enable_rerank": true,
|
||||
"response_type": "Multiple Paragraphs"
|
||||
}' "hybrid"
|
||||
|
||||
# Output test result statistics
|
||||
echo ""
|
||||
echo "=================================================="
|
||||
echo -e "${BLUE}Test Result Statistics${NC}"
|
||||
echo "=================================================="
|
||||
echo -e "Total tests: ${BLUE}$TOTAL_TESTS${NC}"
|
||||
echo -e "Passed tests: ${GREEN}$PASSED_TESTS${NC}"
|
||||
echo -e "Failed tests: ${RED}$FAILED_TESTS${NC}"
|
||||
|
||||
if [[ $FAILED_TESTS -eq 0 ]]; then
|
||||
echo -e "${GREEN}🎉 All tests passed! New data format adaptation successful!${NC}"
|
||||
exit 0
|
||||
else
|
||||
echo -e "${RED}⚠️ $FAILED_TESTS test(s) failed, please check the issues${NC}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "💡 Usage Instructions:"
|
||||
echo "1. Ensure LightRAG API service is running (python -m lightrag.api.lightrag_server)"
|
||||
echo "2. Adjust BASE_URL as needed"
|
||||
echo "3. If authentication is required, add -H \"Authorization: Bearer your-token\""
|
||||
echo "4. Install jq for better JSON formatting output: brew install jq (macOS) or apt install jq (Ubuntu)"
|
||||
echo "5. Script will automatically validate new data format structure: status, message, data, metadata"
|
||||
@@ -0,0 +1,403 @@
|
||||
import pytest
|
||||
|
||||
from lightrag.constants import SOURCE_IDS_LIMIT_METHOD_KEEP
|
||||
from lightrag.constants import GRAPH_FIELD_SEP
|
||||
from lightrag.operate import (
|
||||
_handle_single_entity_extraction,
|
||||
_merge_nodes_then_upsert,
|
||||
_normalize_text_extraction_record_attributes,
|
||||
_handle_single_relationship_extraction,
|
||||
)
|
||||
from lightrag import utils_graph
|
||||
from lightrag.utils import VectorStorageConsistencyError
|
||||
|
||||
|
||||
class DummyGraphStorage:
|
||||
def __init__(self, node=None):
|
||||
self.node = node
|
||||
self.upserted_nodes = []
|
||||
|
||||
async def get_node(self, node_id):
|
||||
return self.node
|
||||
|
||||
async def upsert_node(self, node_id, node_data):
|
||||
self.upserted_nodes.append((node_id, node_data))
|
||||
self.node = dict(node_data)
|
||||
|
||||
|
||||
class DummyVectorStorage:
|
||||
def __init__(self):
|
||||
self.global_config = {"workspace": "test"}
|
||||
self.upserts = []
|
||||
self.deletes = []
|
||||
|
||||
async def upsert(self, data):
|
||||
self.upserts.append(data)
|
||||
return None
|
||||
|
||||
async def delete(self, ids):
|
||||
self.deletes.append(ids)
|
||||
return None
|
||||
|
||||
async def get_by_id(self, id_):
|
||||
return None
|
||||
|
||||
async def index_done_callback(self):
|
||||
return True
|
||||
|
||||
|
||||
class DummyAsyncContext:
|
||||
async def __aenter__(self):
|
||||
return None
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
|
||||
class DummyMergeGraphStorage:
|
||||
def __init__(self):
|
||||
self.nodes = {
|
||||
"Canonical": {
|
||||
"entity_id": "Canonical",
|
||||
"description": "canonical desc",
|
||||
"entity_type": "ORG",
|
||||
"source_id": "chunk-1",
|
||||
"file_path": "canonical.md",
|
||||
},
|
||||
"Alias": {
|
||||
"entity_id": "Alias",
|
||||
"description": "alias desc",
|
||||
"entity_type": "ORG",
|
||||
"source_id": "chunk-2",
|
||||
"file_path": "alias.md",
|
||||
},
|
||||
"Neighbor": {
|
||||
"entity_id": "Neighbor",
|
||||
"description": "neighbor desc",
|
||||
"entity_type": "ORG",
|
||||
"source_id": "chunk-3",
|
||||
"file_path": "neighbor.md",
|
||||
},
|
||||
}
|
||||
self.edges = {
|
||||
("Alias", "Neighbor"): {
|
||||
"description": "rel desc",
|
||||
"keywords": "alias",
|
||||
"source_id": "chunk-rel",
|
||||
"weight": 1.0,
|
||||
"file_path": "rel.md",
|
||||
}
|
||||
}
|
||||
|
||||
async def has_node(self, node_id):
|
||||
return node_id in self.nodes
|
||||
|
||||
async def get_node(self, node_id):
|
||||
return self.nodes[node_id]
|
||||
|
||||
async def upsert_node(self, node_id, node_data):
|
||||
self.nodes[node_id] = dict(node_data)
|
||||
|
||||
async def get_node_edges(self, node_id):
|
||||
results = []
|
||||
for src, tgt in self.edges:
|
||||
if src == node_id or tgt == node_id:
|
||||
results.append((src, tgt))
|
||||
return results
|
||||
|
||||
async def get_edge(self, src, tgt):
|
||||
return self.edges.get((src, tgt)) or self.edges.get((tgt, src))
|
||||
|
||||
async def upsert_edge(self, src, tgt, edge_data):
|
||||
self.edges[(src, tgt)] = dict(edge_data)
|
||||
|
||||
async def delete_node(self, node_id):
|
||||
self.nodes.pop(node_id, None)
|
||||
self.edges = {
|
||||
(src, tgt): data
|
||||
for (src, tgt), data in self.edges.items()
|
||||
if src != node_id and tgt != node_id
|
||||
}
|
||||
|
||||
async def index_done_callback(self):
|
||||
return True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_merge_nodes_then_upsert_handles_missing_legacy_description():
|
||||
graph = DummyGraphStorage(node={"source_id": "chunk-1"})
|
||||
global_config = {
|
||||
"source_ids_limit_method": SOURCE_IDS_LIMIT_METHOD_KEEP,
|
||||
"max_source_ids_per_entity": 20,
|
||||
}
|
||||
|
||||
result = await _merge_nodes_then_upsert(
|
||||
entity_name="LegacyEntity",
|
||||
nodes_data=[],
|
||||
knowledge_graph_inst=graph,
|
||||
entity_vdb=None,
|
||||
global_config=global_config,
|
||||
)
|
||||
|
||||
assert result["description"] == "Entity LegacyEntity"
|
||||
assert graph.upserted_nodes[-1][1]["description"] == "Entity LegacyEntity"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acreate_entity_rejects_empty_description():
|
||||
with pytest.raises(ValueError, match="description cannot be empty"):
|
||||
await utils_graph.acreate_entity(
|
||||
chunk_entity_relation_graph=None,
|
||||
entities_vdb=None,
|
||||
relationships_vdb=None,
|
||||
entity_name="EntityA",
|
||||
entity_data={"description": " "},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acreate_relation_rejects_empty_description():
|
||||
with pytest.raises(ValueError, match="description cannot be empty"):
|
||||
await utils_graph.acreate_relation(
|
||||
chunk_entity_relation_graph=None,
|
||||
entities_vdb=None,
|
||||
relationships_vdb=None,
|
||||
source_entity="A",
|
||||
target_entity="B",
|
||||
relation_data={"description": ""},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aedit_entity_rejects_empty_description():
|
||||
with pytest.raises(ValueError, match="description cannot be empty"):
|
||||
await utils_graph.aedit_entity(
|
||||
chunk_entity_relation_graph=None,
|
||||
entities_vdb=None,
|
||||
relationships_vdb=None,
|
||||
entity_name="EntityA",
|
||||
updated_data={"description": None},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aedit_relation_rejects_empty_description():
|
||||
with pytest.raises(ValueError, match="description cannot be empty"):
|
||||
await utils_graph.aedit_relation(
|
||||
chunk_entity_relation_graph=None,
|
||||
entities_vdb=None,
|
||||
relationships_vdb=None,
|
||||
source_entity="A",
|
||||
target_entity="B",
|
||||
updated_data={"description": " "},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aedit_entity_allows_updates_without_description(monkeypatch):
|
||||
async def fake_edit_impl(*args, **kwargs):
|
||||
return {"entity_name": "EntityA", "description": "kept", "source_id": "chunk-1"}
|
||||
|
||||
monkeypatch.setattr(
|
||||
utils_graph, "get_storage_keyed_lock", lambda *a, **k: DummyAsyncContext()
|
||||
)
|
||||
monkeypatch.setattr(utils_graph, "_edit_entity_impl", fake_edit_impl)
|
||||
|
||||
result = await utils_graph.aedit_entity(
|
||||
chunk_entity_relation_graph=None,
|
||||
entities_vdb=DummyVectorStorage(),
|
||||
relationships_vdb=DummyVectorStorage(),
|
||||
entity_name="EntityA",
|
||||
updated_data={"entity_type": "ORG"},
|
||||
)
|
||||
|
||||
assert result["operation_summary"]["operation_status"] == "success"
|
||||
|
||||
|
||||
def test_handle_single_relationship_extraction_ignores_empty_description():
|
||||
relation = _handle_single_relationship_extraction(
|
||||
["relation", "Alice", "Bob", "works_with", " "],
|
||||
chunk_key="chunk-1",
|
||||
timestamp=1,
|
||||
)
|
||||
|
||||
assert relation is None
|
||||
|
||||
|
||||
def test_mis_prefixed_relation_row_is_recovered():
|
||||
record = _normalize_text_extraction_record_attributes(
|
||||
["entity", "Alice", "Acme Corp", "founded", "Alice founded Acme Corp."],
|
||||
chunk_key="chunk-1",
|
||||
)
|
||||
|
||||
relation = _handle_single_relationship_extraction(
|
||||
record,
|
||||
chunk_key="chunk-1",
|
||||
timestamp=1,
|
||||
)
|
||||
|
||||
assert relation is not None
|
||||
assert relation["src_id"] == "Alice"
|
||||
assert relation["tgt_id"] == "Acme Corp"
|
||||
|
||||
|
||||
def test_four_part_entity_row_remains_entity():
|
||||
record = _normalize_text_extraction_record_attributes(
|
||||
["entity", "Alice", "Person", "Alice is the founder of Acme Corp."],
|
||||
chunk_key="chunk-1",
|
||||
)
|
||||
|
||||
entity = _handle_single_entity_extraction(
|
||||
record,
|
||||
chunk_key="chunk-1",
|
||||
timestamp=1,
|
||||
)
|
||||
|
||||
assert entity is not None
|
||||
assert entity["entity_name"] == "Alice"
|
||||
|
||||
|
||||
def test_malformed_recovered_relation_still_fails():
|
||||
record = _normalize_text_extraction_record_attributes(
|
||||
["entity", "Alice", "Acme Corp", "founded", " "],
|
||||
chunk_key="chunk-1",
|
||||
)
|
||||
|
||||
relation = _handle_single_relationship_extraction(
|
||||
record,
|
||||
chunk_key="chunk-1",
|
||||
timestamp=1,
|
||||
)
|
||||
|
||||
assert relation is None
|
||||
|
||||
|
||||
def test_unrelated_five_part_prefix_remains_invalid():
|
||||
record = _normalize_text_extraction_record_attributes(
|
||||
["edge", "Alice", "Acme Corp", "founded", "Alice founded Acme Corp."],
|
||||
chunk_key="chunk-1",
|
||||
)
|
||||
|
||||
relation = _handle_single_relationship_extraction(
|
||||
record,
|
||||
chunk_key="chunk-1",
|
||||
timestamp=1,
|
||||
)
|
||||
|
||||
assert relation is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_merge_entities_preserves_file_path_in_vector_updates(monkeypatch):
|
||||
graph = DummyMergeGraphStorage()
|
||||
entities_vdb = DummyVectorStorage()
|
||||
relationships_vdb = DummyVectorStorage()
|
||||
|
||||
async def fake_get_entity_info(*args, **kwargs):
|
||||
return {"entity_name": "Canonical"}
|
||||
|
||||
monkeypatch.setattr(utils_graph, "get_entity_info", fake_get_entity_info)
|
||||
|
||||
await utils_graph._merge_entities_impl(
|
||||
chunk_entity_relation_graph=graph,
|
||||
entities_vdb=entities_vdb,
|
||||
relationships_vdb=relationships_vdb,
|
||||
source_entities=["Alias", "Canonical"],
|
||||
target_entity="Canonical",
|
||||
)
|
||||
|
||||
relationship_payload = relationships_vdb.upserts[-1]
|
||||
entity_payload = entities_vdb.upserts[-1]
|
||||
|
||||
assert next(iter(relationship_payload.values()))["file_path"] == "rel.md"
|
||||
assert set(
|
||||
next(iter(entity_payload.values()))["file_path"].split(GRAPH_FIELD_SEP)
|
||||
) == {
|
||||
"alias.md",
|
||||
"canonical.md",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_merge_response_is_built_from_graph_without_reading_vdb():
|
||||
# The merge response is built from graph data only. Reading the vector
|
||||
# store is redundant (the graph is authoritative) and previously leaked a
|
||||
# non-JSON-serializable embedding (pgvector -> numpy) into the API
|
||||
# response, causing a 500 on /graph/entity/edit with all-PostgreSQL.
|
||||
graph = DummyMergeGraphStorage()
|
||||
entities_vdb = DummyVectorStorage()
|
||||
relationships_vdb = DummyVectorStorage()
|
||||
|
||||
vdb_reads = []
|
||||
inner_get_by_id = entities_vdb.get_by_id
|
||||
|
||||
async def _tracked_get_by_id(id_):
|
||||
vdb_reads.append(id_)
|
||||
return await inner_get_by_id(id_)
|
||||
|
||||
entities_vdb.get_by_id = _tracked_get_by_id
|
||||
|
||||
result = await utils_graph._merge_entities_impl(
|
||||
chunk_entity_relation_graph=graph,
|
||||
entities_vdb=entities_vdb,
|
||||
relationships_vdb=relationships_vdb,
|
||||
source_entities=["Alias", "Canonical"],
|
||||
target_entity="Canonical",
|
||||
)
|
||||
|
||||
assert "graph_data" in result
|
||||
assert "vector_data" not in result
|
||||
# The response did not read the entity vector store.
|
||||
assert vdb_reads == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aedit_entity_merge_propagates_consistency_error(monkeypatch):
|
||||
"""rename-with-merge must surface VectorStorageConsistencyError, not swallow it.
|
||||
|
||||
The graph is updated during the merge, then the (deferred) vector-store
|
||||
flush fails. aedit_entity used to fold every merge exception into a
|
||||
partial-success summary and return normally, so the /graph/entity/edit
|
||||
route reported HTTP 200 "success" and hid the consistency failure that the
|
||||
fail-loud merge path exists to surface. It must re-raise instead.
|
||||
"""
|
||||
graph = DummyMergeGraphStorage()
|
||||
entities_vdb = DummyVectorStorage()
|
||||
relationships_vdb = DummyVectorStorage()
|
||||
|
||||
# Deferred-embedding flush failure: upsert buffers, index_done_callback raises.
|
||||
async def _boom():
|
||||
raise RuntimeError("embedder down")
|
||||
|
||||
relationships_vdb.index_done_callback = _boom
|
||||
|
||||
# Single-attempt VDB wrapper (no retry delays) + a no-op storage lock.
|
||||
async def _single_attempt(operation, **kwargs):
|
||||
await operation()
|
||||
|
||||
monkeypatch.setattr(
|
||||
utils_graph, "safe_vdb_operation_with_exception", _single_attempt
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
utils_graph, "get_storage_keyed_lock", lambda *a, **k: DummyAsyncContext()
|
||||
)
|
||||
|
||||
async def fake_get_entity_info(*args, **kwargs):
|
||||
return {"entity_name": "Canonical"}
|
||||
|
||||
monkeypatch.setattr(utils_graph, "get_entity_info", fake_get_entity_info)
|
||||
|
||||
with pytest.raises(VectorStorageConsistencyError) as excinfo:
|
||||
await utils_graph.aedit_entity(
|
||||
chunk_entity_relation_graph=graph,
|
||||
entities_vdb=entities_vdb,
|
||||
relationships_vdb=relationships_vdb,
|
||||
entity_name="Alias",
|
||||
updated_data={"entity_name": "Canonical"},
|
||||
allow_rename=True,
|
||||
allow_merge=True,
|
||||
)
|
||||
|
||||
assert "lightrag-rebuild-vdb" in str(excinfo.value)
|
||||
# Fail-loud happened before source deletion: 'Alias' is still in the graph.
|
||||
assert "Alias" in graph.nodes
|
||||
@@ -0,0 +1,694 @@
|
||||
"""Tests for the `/documents/text(s)` ``chunking`` request object.
|
||||
|
||||
Three concerns:
|
||||
|
||||
1. **Synchronous validation**: malformed ``chunking`` is rejected at
|
||||
request-parse time (HTTP 422 / ``ValidationError``) — never deferred to
|
||||
the background indexing task, where the HTTP response is already sent.
|
||||
The per-strategy typed params models do full type + value checking, not
|
||||
just unknown-key detection.
|
||||
|
||||
2. **``_resolve_text_chunking``**: a validated ``chunking`` config is frozen
|
||||
into ``(process_options, chunk_options)``; ``chunk_token_size`` and the
|
||||
strategy params land in the selected strategy's sub-dict, overriding any
|
||||
env-derived value, while the other strategy sub-dicts are dropped (slim).
|
||||
|
||||
3. **Route forwarding**: ``/documents/text`` and ``/documents/texts`` forward
|
||||
``request.chunking`` to ``pipeline_index_texts`` and return 422 (without
|
||||
scheduling any background work) for a malformed body.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
from pydantic import ValidationError
|
||||
|
||||
_original_argv = sys.argv[:]
|
||||
sys.argv = [sys.argv[0]]
|
||||
_dr = importlib.import_module("lightrag.api.routers.document_routes")
|
||||
sys.argv = _original_argv
|
||||
|
||||
TextChunkingConfig = _dr.TextChunkingConfig
|
||||
InsertTextRequest = _dr.InsertTextRequest
|
||||
_resolve_text_chunking = _dr._resolve_text_chunking
|
||||
create_document_routes = _dr.create_document_routes
|
||||
|
||||
from lightrag.constants import ( # noqa: E402
|
||||
PROCESS_OPTION_CHUNK_FIXED,
|
||||
PROCESS_OPTION_CHUNK_PARAGRAH,
|
||||
PROCESS_OPTION_CHUNK_RECURSIVE,
|
||||
PROCESS_OPTION_CHUNK_VECTOR,
|
||||
)
|
||||
from lightrag.parser.routing import default_chunker_config # noqa: E402
|
||||
|
||||
pytestmark = pytest.mark.offline
|
||||
|
||||
_ALL_STRATEGY_KEYS = {
|
||||
"fixed_token",
|
||||
"recursive_character",
|
||||
"semantic_vector",
|
||||
"paragraph_semantic",
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. Synchronous validation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
# wrong types (strict rejects lax coercion)
|
||||
{"strategy": "fixed_token", "params": {"chunk_token_size": True}},
|
||||
{"strategy": "fixed_token", "params": {"chunk_token_size": "5"}},
|
||||
{"strategy": "fixed_token", "params": {"chunk_token_size": 1.5}},
|
||||
{"strategy": "fixed_token", "params": {"chunk_overlap_token_size": "bad"}},
|
||||
{"strategy": "fixed_token", "params": {"split_by_character": 123}},
|
||||
{"strategy": "fixed_token", "params": {"split_by_character_only": 1}},
|
||||
{"strategy": "recursive_character", "params": {"separators": "abc"}},
|
||||
{"strategy": "recursive_character", "params": {"separators": [1, 2]}},
|
||||
# value / range
|
||||
{"strategy": "fixed_token", "params": {"chunk_token_size": 0}},
|
||||
{"strategy": "recursive_character", "params": {"chunk_overlap_token_size": -1}},
|
||||
{"strategy": "semantic_vector", "params": {"buffer_size": 0}},
|
||||
{"strategy": "semantic_vector", "params": {"buffer_size": True}},
|
||||
{
|
||||
"strategy": "semantic_vector",
|
||||
"params": {"breakpoint_threshold_type": "p99"},
|
||||
},
|
||||
{
|
||||
"strategy": "semantic_vector",
|
||||
"params": {"breakpoint_threshold_amount": 0},
|
||||
},
|
||||
{
|
||||
# strict float rejects strings (no lax numeric-string coercion)
|
||||
"strategy": "semantic_vector",
|
||||
"params": {"breakpoint_threshold_amount": "95"},
|
||||
},
|
||||
{
|
||||
# strict float rejects bool (bool is an int subclass, undesirable here)
|
||||
"strategy": "semantic_vector",
|
||||
"params": {"breakpoint_threshold_amount": True},
|
||||
},
|
||||
{
|
||||
# > 100 with an explicit percentile/gradient type is rejected at
|
||||
# parse time (both fields present, no inheritance ambiguity).
|
||||
"strategy": "semantic_vector",
|
||||
"params": {
|
||||
"breakpoint_threshold_type": "percentile",
|
||||
"breakpoint_threshold_amount": 150,
|
||||
},
|
||||
},
|
||||
{
|
||||
# malformed regex must be compiled/rejected at parse time
|
||||
"strategy": "semantic_vector",
|
||||
"params": {"sentence_split_regex": "("},
|
||||
},
|
||||
# cross-field
|
||||
{
|
||||
"strategy": "fixed_token",
|
||||
"params": {"chunk_token_size": 100, "chunk_overlap_token_size": 200},
|
||||
},
|
||||
# unknown / wrong-for-strategy keys
|
||||
{"strategy": "fixed_token", "params": {"bogus": 1}},
|
||||
{"strategy": "fixed_token", "params": {"separators": ["x"]}},
|
||||
{"strategy": "recursive_character", "params": {"buffer_size": 1}},
|
||||
],
|
||||
)
|
||||
def test_chunking_config_rejects_malformed(body):
|
||||
with pytest.raises(ValidationError):
|
||||
TextChunkingConfig.model_validate(body)
|
||||
|
||||
|
||||
def test_chunking_config_defaults_to_fixed_token():
|
||||
cfg = TextChunkingConfig.model_validate({"params": {"chunk_token_size": 500}})
|
||||
assert cfg.strategy == "fixed_token"
|
||||
assert cfg.params == {"chunk_token_size": 500}
|
||||
|
||||
|
||||
def test_chunking_config_normalizes_to_supplied_keys_only():
|
||||
# int amount is coerced to float; only the supplied key survives.
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{"strategy": "semantic_vector", "params": {"breakpoint_threshold_amount": 95}}
|
||||
)
|
||||
assert cfg.params == {"breakpoint_threshold_amount": 95.0}
|
||||
|
||||
|
||||
def test_chunking_config_amount_in_range_for_std_deviation():
|
||||
# standard_deviation only requires > 0 (no [0, 100] ceiling).
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{
|
||||
"strategy": "semantic_vector",
|
||||
"params": {
|
||||
"breakpoint_threshold_type": "standard_deviation",
|
||||
"breakpoint_threshold_amount": 3.5,
|
||||
},
|
||||
}
|
||||
)
|
||||
assert cfg.params["breakpoint_threshold_amount"] == 3.5
|
||||
|
||||
|
||||
def test_chunking_config_amount_over_100_without_type_is_deferred():
|
||||
# Type omitted -> the (0, 100] ceiling cannot be decided at parse time
|
||||
# (the effective type may be inherited), so the model must NOT assume
|
||||
# percentile and reject. _resolve_text_chunking applies the ceiling later.
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{"strategy": "semantic_vector", "params": {"breakpoint_threshold_amount": 150}}
|
||||
)
|
||||
assert cfg.params == {"breakpoint_threshold_amount": 150.0}
|
||||
|
||||
|
||||
def test_chunking_config_accepts_int_amount_widened_to_float():
|
||||
# Strict float accepts an int (JSON 95) and widens it to 95.0 — the common
|
||||
# documented threshold magnitude. (str/bool are rejected; see the
|
||||
# rejection matrix above.) Exercised via both python and JSON validation
|
||||
# modes so the FastAPI request path (which parses JSON) stays covered.
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{"strategy": "semantic_vector", "params": {"breakpoint_threshold_amount": 95}}
|
||||
)
|
||||
assert cfg.params == {"breakpoint_threshold_amount": 95.0}
|
||||
assert isinstance(cfg.params["breakpoint_threshold_amount"], float)
|
||||
|
||||
cfg_json = TextChunkingConfig.model_validate_json(
|
||||
'{"strategy": "semantic_vector", "params": {"breakpoint_threshold_amount": 95}}'
|
||||
)
|
||||
assert cfg_json.params == {"breakpoint_threshold_amount": 95.0}
|
||||
|
||||
|
||||
def test_chunking_config_accepts_valid_sentence_split_regex():
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{
|
||||
"strategy": "semantic_vector",
|
||||
"params": {"sentence_split_regex": r"(?<=[.?!])\s+"},
|
||||
}
|
||||
)
|
||||
assert cfg.params == {"sentence_split_regex": r"(?<=[.?!])\s+"}
|
||||
|
||||
|
||||
def test_chunking_config_drops_explicit_null():
|
||||
# Explicit null means "inherit the default" (every param field is
|
||||
# Optional/None=inherit), so it must be dropped — not merged over the
|
||||
# resolved default, which would later make the chunker do int(None).
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{"strategy": "fixed_token", "params": {"chunk_token_size": None}}
|
||||
)
|
||||
assert cfg.params == {}
|
||||
|
||||
|
||||
def test_chunking_config_keeps_real_value_drops_sibling_null():
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{
|
||||
"strategy": "fixed_token",
|
||||
"params": {"chunk_token_size": 500, "split_by_character": None},
|
||||
}
|
||||
)
|
||||
assert cfg.params == {"chunk_token_size": 500}
|
||||
|
||||
|
||||
def test_insert_text_request_rejects_malformed_chunking():
|
||||
with pytest.raises(ValidationError):
|
||||
InsertTextRequest.model_validate(
|
||||
{
|
||||
"text": "hi",
|
||||
"file_source": "a.md",
|
||||
"chunking": {
|
||||
"strategy": "recursive_character",
|
||||
"params": {"separators": "notalist"},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. _resolve_text_chunking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _stub_rag(addon_params=None):
|
||||
return SimpleNamespace(
|
||||
addon_params=addon_params if addon_params is not None else {}
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_none_keeps_default_fixed():
|
||||
process_options, chunk_options = _resolve_text_chunking(None, _stub_rag())
|
||||
assert process_options == PROCESS_OPTION_CHUNK_FIXED
|
||||
assert "fixed_token" in chunk_options
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"strategy,expected_po,key",
|
||||
[
|
||||
("fixed_token", PROCESS_OPTION_CHUNK_FIXED, "fixed_token"),
|
||||
("recursive_character", PROCESS_OPTION_CHUNK_RECURSIVE, "recursive_character"),
|
||||
("semantic_vector", PROCESS_OPTION_CHUNK_VECTOR, "semantic_vector"),
|
||||
("paragraph_semantic", PROCESS_OPTION_CHUNK_PARAGRAH, "paragraph_semantic"),
|
||||
],
|
||||
)
|
||||
def test_resolve_maps_strategy_and_writes_size_into_subdict(strategy, expected_po, key):
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{"strategy": strategy, "params": {"chunk_token_size": 777}}
|
||||
)
|
||||
process_options, chunk_options = _resolve_text_chunking(cfg, _stub_rag())
|
||||
assert process_options == expected_po
|
||||
# chunk_token_size lands in the strategy sub-dict for ALL strategies
|
||||
# (F included, post-cleanup) — that's where process_single_document reads it.
|
||||
assert chunk_options[key]["chunk_token_size"] == 777
|
||||
# slim contract: other strategies' sub-dicts are dropped
|
||||
for other in _ALL_STRATEGY_KEYS - {key}:
|
||||
assert other not in chunk_options
|
||||
|
||||
|
||||
def test_resolve_merges_strategy_params():
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{
|
||||
"strategy": "recursive_character",
|
||||
"params": {"separators": ["A", "B"], "chunk_overlap_token_size": 0},
|
||||
}
|
||||
)
|
||||
_, chunk_options = _resolve_text_chunking(cfg, _stub_rag())
|
||||
assert chunk_options["recursive_character"]["separators"] == ["A", "B"]
|
||||
assert chunk_options["recursive_character"]["chunk_overlap_token_size"] == 0
|
||||
|
||||
|
||||
def test_resolve_size_overrides_env_for_recursive(monkeypatch):
|
||||
monkeypatch.setenv("CHUNK_R_SIZE", "999")
|
||||
addon = {"chunker": default_chunker_config()}
|
||||
# sanity: env baked into the R sub-dict
|
||||
assert addon["chunker"]["recursive_character"]["chunk_token_size"] == 999
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{"strategy": "recursive_character", "params": {"chunk_token_size": 1234}}
|
||||
)
|
||||
_, chunk_options = _resolve_text_chunking(cfg, _stub_rag(addon))
|
||||
# API value wins over the env-derived sub-dict value.
|
||||
assert chunk_options["recursive_character"]["chunk_token_size"] == 1234
|
||||
|
||||
|
||||
def test_resolve_split_by_character_only_false_overrides_env(monkeypatch):
|
||||
# The API path can express an explicit False (a plain dict merge), unlike
|
||||
# the ainsert positional-arg path. Prove it overrides an env-True default.
|
||||
monkeypatch.setenv("CHUNK_F_SPLIT_BY_CHARACTER_ONLY", "true")
|
||||
addon = {"chunker": default_chunker_config()}
|
||||
assert addon["chunker"]["fixed_token"]["split_by_character_only"] is True
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{"strategy": "fixed_token", "params": {"split_by_character_only": False}}
|
||||
)
|
||||
_, chunk_options = _resolve_text_chunking(cfg, _stub_rag(addon))
|
||||
assert chunk_options["fixed_token"]["split_by_character_only"] is False
|
||||
|
||||
|
||||
def test_resolve_drop_references_request_false_overrides_env_true(monkeypatch):
|
||||
# The JSON text API can express an explicit ``drop_references=False`` that
|
||||
# overrides ``CHUNK_P_DROP_REFERENCES=true`` (resolved via slim_chunk_options
|
||||
# then overlaid by the request params). Mirrors the split_by_character_only
|
||||
# override above for the paragraph-semantic switch.
|
||||
monkeypatch.setenv("CHUNK_P_DROP_REFERENCES", "true")
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{"strategy": "paragraph_semantic", "params": {"drop_references": False}}
|
||||
)
|
||||
_, chunk_options = _resolve_text_chunking(cfg, _stub_rag())
|
||||
assert chunk_options["paragraph_semantic"]["drop_references"] is False
|
||||
|
||||
|
||||
def test_resolve_drop_references_env_true_without_request_param(monkeypatch):
|
||||
# No request param → the env-true default survives into the snapshot.
|
||||
monkeypatch.setenv("CHUNK_P_DROP_REFERENCES", "true")
|
||||
cfg = TextChunkingConfig.model_validate({"strategy": "paragraph_semantic"})
|
||||
_, chunk_options = _resolve_text_chunking(cfg, _stub_rag())
|
||||
assert chunk_options["paragraph_semantic"]["drop_references"] is True
|
||||
|
||||
|
||||
def test_resolve_rejects_size_below_inherited_overlap(monkeypatch):
|
||||
# Overlap is inherited from addon_params (not in the request), so the
|
||||
# request model can't catch it — _resolve_text_chunking must.
|
||||
monkeypatch.setenv("CHUNK_F_OVERLAP_SIZE", "100")
|
||||
addon = {"chunker": default_chunker_config()}
|
||||
assert addon["chunker"]["fixed_token"]["chunk_overlap_token_size"] == 100
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{"strategy": "fixed_token", "params": {"chunk_token_size": 50}}
|
||||
)
|
||||
with pytest.raises(ValueError, match="chunk_overlap_token_size"):
|
||||
_resolve_text_chunking(cfg, _stub_rag(addon))
|
||||
|
||||
|
||||
def test_resolve_allows_size_above_inherited_overlap(monkeypatch):
|
||||
monkeypatch.setenv("CHUNK_F_OVERLAP_SIZE", "100")
|
||||
addon = {"chunker": default_chunker_config()}
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{"strategy": "fixed_token", "params": {"chunk_token_size": 400}}
|
||||
)
|
||||
_, chunk_options = _resolve_text_chunking(cfg, _stub_rag(addon))
|
||||
assert chunk_options["fixed_token"]["chunk_token_size"] == 400
|
||||
|
||||
|
||||
def test_resolve_skips_overlap_check_for_delimiter_only(monkeypatch):
|
||||
# Delimiter-only fixed-token chunking never uses overlap, so a small
|
||||
# chunk_token_size below the inherited overlap must NOT be rejected.
|
||||
monkeypatch.setenv("CHUNK_F_OVERLAP_SIZE", "100")
|
||||
addon = {"chunker": default_chunker_config()}
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{
|
||||
"strategy": "fixed_token",
|
||||
"params": {
|
||||
"split_by_character": "\n\n",
|
||||
"split_by_character_only": True,
|
||||
"chunk_token_size": 50,
|
||||
},
|
||||
}
|
||||
)
|
||||
_, chunk_options = _resolve_text_chunking(cfg, _stub_rag(addon))
|
||||
assert chunk_options["fixed_token"]["chunk_token_size"] == 50
|
||||
|
||||
|
||||
def test_resolve_enforces_overlap_when_only_flag_without_delimiter(monkeypatch):
|
||||
# split_by_character_only is a no-op without split_by_character: the chunker
|
||||
# falls back to normal token windowing, which DOES use overlap — so the
|
||||
# overlap < size check must still fire here.
|
||||
monkeypatch.setenv("CHUNK_F_OVERLAP_SIZE", "100")
|
||||
monkeypatch.delenv("CHUNK_F_SPLIT_BY_CHARACTER", raising=False)
|
||||
addon = {"chunker": default_chunker_config()}
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{
|
||||
"strategy": "fixed_token",
|
||||
"params": {"split_by_character_only": True, "chunk_token_size": 50},
|
||||
}
|
||||
)
|
||||
with pytest.raises(ValueError, match="chunk_overlap_token_size"):
|
||||
_resolve_text_chunking(cfg, _stub_rag(addon))
|
||||
|
||||
|
||||
def test_resolve_allows_amount_over_100_with_inherited_std_type():
|
||||
# Request overrides only the amount; the standard_deviation type is
|
||||
# inherited from addon_params. std/iqr have no (0, 100] ceiling, so this
|
||||
# must NOT be rejected (the request model deferred the check here).
|
||||
addon = {
|
||||
"chunker": {
|
||||
"semantic_vector": {"breakpoint_threshold_type": "standard_deviation"}
|
||||
}
|
||||
}
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{"strategy": "semantic_vector", "params": {"breakpoint_threshold_amount": 150}}
|
||||
)
|
||||
_, chunk_options = _resolve_text_chunking(cfg, _stub_rag(addon))
|
||||
assert chunk_options["semantic_vector"]["breakpoint_threshold_amount"] == 150
|
||||
assert (
|
||||
chunk_options["semantic_vector"]["breakpoint_threshold_type"]
|
||||
== "standard_deviation"
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_rejects_amount_over_100_with_inherited_percentile_type():
|
||||
# Same partial override, but the effective (inherited) type is percentile,
|
||||
# which feeds np.percentile and requires the (0, 100] ceiling.
|
||||
addon = {
|
||||
"chunker": {"semantic_vector": {"breakpoint_threshold_type": "percentile"}}
|
||||
}
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{"strategy": "semantic_vector", "params": {"breakpoint_threshold_amount": 150}}
|
||||
)
|
||||
with pytest.raises(ValueError, match="breakpoint_threshold_amount"):
|
||||
_resolve_text_chunking(cfg, _stub_rag(addon))
|
||||
|
||||
|
||||
def test_resolve_null_size_does_not_erase_inherited_default(monkeypatch):
|
||||
# An explicit null in the request must not overwrite the resolved size
|
||||
# with None (which would make the chunker do int(None) in the background).
|
||||
monkeypatch.setenv("CHUNK_F_SIZE", "640")
|
||||
addon = {"chunker": default_chunker_config()}
|
||||
cfg = TextChunkingConfig.model_validate(
|
||||
{"strategy": "fixed_token", "params": {"chunk_token_size": None}}
|
||||
)
|
||||
_, chunk_options = _resolve_text_chunking(cfg, _stub_rag(addon))
|
||||
# null dropped by the model -> inherited CHUNK_F_SIZE survives, no None.
|
||||
assert chunk_options["fixed_token"]["chunk_token_size"] == 640
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. Route forwarding + synchronous 422
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FwdDocStatus:
|
||||
async def get_doc_by_file_basename(self, basename):
|
||||
return None
|
||||
|
||||
|
||||
class _FwdRag:
|
||||
workspace = "chunk-fwd-test"
|
||||
addon_params: dict = {}
|
||||
|
||||
def __init__(self):
|
||||
self.doc_status = _FwdDocStatus()
|
||||
|
||||
|
||||
_HEADERS = {"X-API-Key": "test-key"}
|
||||
|
||||
|
||||
def _make_client(monkeypatch, addon_params=None):
|
||||
"""Build a TestClient whose enqueue-slot guards are no-ops and whose
|
||||
``pipeline_index_texts`` is a spy recording the forwarded args.
|
||||
|
||||
``addon_params`` seeds the rag the routes resolve chunking against; the
|
||||
handler calls the real ``_resolve_text_chunking`` synchronously, so the
|
||||
effective-overlap validation runs against this snapshot.
|
||||
"""
|
||||
captured: dict = {}
|
||||
|
||||
async def _spy(rag, texts, file_sources=None, track_id=None, chunking=None):
|
||||
captured["texts"] = texts
|
||||
captured["file_sources"] = file_sources
|
||||
captured["chunking"] = chunking
|
||||
|
||||
async def _noop_reserve(rag):
|
||||
return False
|
||||
|
||||
async def _noop_release(rag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(_dr, "pipeline_index_texts", _spy)
|
||||
monkeypatch.setattr(_dr, "_reserve_enqueue_slot", _noop_reserve)
|
||||
monkeypatch.setattr(_dr, "_release_enqueue_slot", _noop_release)
|
||||
|
||||
rag = _FwdRag()
|
||||
rag.addon_params = addon_params if addon_params is not None else {}
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(
|
||||
create_document_routes(rag, SimpleNamespace(), api_key="test-key")
|
||||
)
|
||||
return TestClient(app), captured
|
||||
|
||||
|
||||
def test_insert_text_forwards_chunking(monkeypatch):
|
||||
client, captured = _make_client(monkeypatch)
|
||||
resp = client.post(
|
||||
"/documents/text",
|
||||
headers=_HEADERS,
|
||||
json={
|
||||
"text": "hello world",
|
||||
"file_source": "a.md",
|
||||
"chunking": {
|
||||
"strategy": "recursive_character",
|
||||
"params": {"chunk_token_size": 1000, "separators": ["X"]},
|
||||
},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert captured["chunking"] is not None
|
||||
assert captured["chunking"].strategy == "recursive_character"
|
||||
assert captured["chunking"].params == {
|
||||
"chunk_token_size": 1000,
|
||||
"separators": ["X"],
|
||||
}
|
||||
|
||||
|
||||
def test_insert_texts_forwards_chunking(monkeypatch):
|
||||
client, captured = _make_client(monkeypatch)
|
||||
resp = client.post(
|
||||
"/documents/texts",
|
||||
headers=_HEADERS,
|
||||
json={
|
||||
"texts": ["one", "two"],
|
||||
"file_sources": ["a.md", "b.md"],
|
||||
"chunking": {"strategy": "semantic_vector", "params": {"buffer_size": 2}},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert captured["chunking"].strategy == "semantic_vector"
|
||||
assert captured["chunking"].params == {"buffer_size": 2}
|
||||
|
||||
|
||||
def test_insert_text_without_chunking_forwards_none(monkeypatch):
|
||||
client, captured = _make_client(monkeypatch)
|
||||
resp = client.post(
|
||||
"/documents/text",
|
||||
headers=_HEADERS,
|
||||
json={"text": "hello", "file_source": "a.md"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert captured["chunking"] is None
|
||||
|
||||
|
||||
def test_insert_text_returns_422_on_malformed_chunking_without_scheduling(monkeypatch):
|
||||
client, captured = _make_client(monkeypatch)
|
||||
resp = client.post(
|
||||
"/documents/text",
|
||||
headers=_HEADERS,
|
||||
json={
|
||||
"text": "hello",
|
||||
"file_source": "a.md",
|
||||
"chunking": {
|
||||
"strategy": "recursive_character",
|
||||
"params": {"separators": "notalist"},
|
||||
},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 422
|
||||
# Body validation fails before the endpoint body runs: no background
|
||||
# indexing is scheduled, so the spy never fires.
|
||||
assert captured == {}
|
||||
|
||||
|
||||
def test_insert_text_returns_422_when_size_below_inherited_overlap(monkeypatch):
|
||||
# chunk_token_size=50 in the request, overlap=100 inherited from the
|
||||
# rag's addon_params (not in the request). The model can't catch this;
|
||||
# the handler's synchronous _resolve_text_chunking must, BEFORE any
|
||||
# background work is scheduled.
|
||||
addon = {
|
||||
"chunker": {
|
||||
"chunk_token_size": 1200,
|
||||
"fixed_token": {"chunk_overlap_token_size": 100},
|
||||
}
|
||||
}
|
||||
client, captured = _make_client(monkeypatch, addon_params=addon)
|
||||
resp = client.post(
|
||||
"/documents/text",
|
||||
headers=_HEADERS,
|
||||
json={
|
||||
"text": "hello",
|
||||
"file_source": "a.md",
|
||||
"chunking": {"strategy": "fixed_token", "params": {"chunk_token_size": 50}},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 422
|
||||
assert "chunk_overlap_token_size" in resp.json()["detail"]
|
||||
# Rejected synchronously: background indexing never scheduled.
|
||||
assert captured == {}
|
||||
|
||||
|
||||
def test_insert_text_allows_amount_override_inheriting_std_type(monkeypatch):
|
||||
# Reviewer scenario: deployment sets standard_deviation; a request
|
||||
# overrides only breakpoint_threshold_amount (> 100). This must be
|
||||
# accepted (not 422), since std has no (0, 100] ceiling.
|
||||
addon = {
|
||||
"chunker": {
|
||||
"semantic_vector": {"breakpoint_threshold_type": "standard_deviation"}
|
||||
}
|
||||
}
|
||||
client, captured = _make_client(monkeypatch, addon_params=addon)
|
||||
resp = client.post(
|
||||
"/documents/text",
|
||||
headers=_HEADERS,
|
||||
json={
|
||||
"text": "hello",
|
||||
"file_source": "a.md",
|
||||
"chunking": {
|
||||
"strategy": "semantic_vector",
|
||||
"params": {"breakpoint_threshold_amount": 150},
|
||||
},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert captured["chunking"].params == {"breakpoint_threshold_amount": 150.0}
|
||||
|
||||
|
||||
def test_insert_text_rejects_amount_over_100_inheriting_percentile_type(monkeypatch):
|
||||
addon = {
|
||||
"chunker": {"semantic_vector": {"breakpoint_threshold_type": "percentile"}}
|
||||
}
|
||||
client, captured = _make_client(monkeypatch, addon_params=addon)
|
||||
resp = client.post(
|
||||
"/documents/text",
|
||||
headers=_HEADERS,
|
||||
json={
|
||||
"text": "hello",
|
||||
"file_source": "a.md",
|
||||
"chunking": {
|
||||
"strategy": "semantic_vector",
|
||||
"params": {"breakpoint_threshold_amount": 150},
|
||||
},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 422
|
||||
assert "breakpoint_threshold_amount" in resp.json()["detail"]
|
||||
assert captured == {}
|
||||
|
||||
|
||||
def test_insert_text_rejects_malformed_sentence_split_regex(monkeypatch):
|
||||
# Malformed regex must 422 at request parse time, before scheduling.
|
||||
client, captured = _make_client(monkeypatch)
|
||||
resp = client.post(
|
||||
"/documents/text",
|
||||
headers=_HEADERS,
|
||||
json={
|
||||
"text": "hello",
|
||||
"file_source": "a.md",
|
||||
"chunking": {
|
||||
"strategy": "semantic_vector",
|
||||
"params": {"sentence_split_regex": "("},
|
||||
},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 422
|
||||
assert captured == {}
|
||||
|
||||
|
||||
def test_insert_text_drops_explicit_null_param(monkeypatch):
|
||||
# "chunk_token_size": null must be treated as "inherit" (dropped), so the
|
||||
# request succeeds and the forwarded params carry no None that would later
|
||||
# crash the chunker with int(None).
|
||||
client, captured = _make_client(monkeypatch)
|
||||
resp = client.post(
|
||||
"/documents/text",
|
||||
headers=_HEADERS,
|
||||
json={
|
||||
"text": "hello",
|
||||
"file_source": "a.md",
|
||||
"chunking": {
|
||||
"strategy": "fixed_token",
|
||||
"params": {"chunk_token_size": None, "chunk_overlap_token_size": 50},
|
||||
},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert captured["chunking"].params == {"chunk_overlap_token_size": 50}
|
||||
|
||||
|
||||
def test_insert_text_allows_small_size_for_delimiter_only(monkeypatch):
|
||||
# Paragraph splitting with a small chunk_token_size: overlap is inherited
|
||||
# (100) but unused in delimiter-only mode, so this must succeed, not 422.
|
||||
addon = {"chunker": {"fixed_token": {"chunk_overlap_token_size": 100}}}
|
||||
client, captured = _make_client(monkeypatch, addon_params=addon)
|
||||
resp = client.post(
|
||||
"/documents/text",
|
||||
headers=_HEADERS,
|
||||
json={
|
||||
"text": "hello",
|
||||
"file_source": "a.md",
|
||||
"chunking": {
|
||||
"strategy": "fixed_token",
|
||||
"params": {
|
||||
"split_by_character": "\n\n",
|
||||
"split_by_character_only": True,
|
||||
"chunk_token_size": 50,
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert captured["chunking"].params["chunk_token_size"] == 50
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,119 @@
|
||||
import importlib
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
_original_argv = sys.argv[:]
|
||||
sys.argv = [sys.argv[0]]
|
||||
_document_routes = importlib.import_module("lightrag.api.routers.document_routes")
|
||||
_base = importlib.import_module("lightrag.base")
|
||||
sys.argv = _original_argv
|
||||
|
||||
create_document_routes = _document_routes.create_document_routes
|
||||
DocProcessingStatus = _base.DocProcessingStatus
|
||||
DocStatus = _base.DocStatus
|
||||
DocStatusStorage = _base.DocStatusStorage
|
||||
|
||||
pytestmark = pytest.mark.offline
|
||||
|
||||
|
||||
def _doc(status: DocStatus, suffix: str) -> DocProcessingStatus:
|
||||
return DocProcessingStatus(
|
||||
content_summary=f"{status.value} summary",
|
||||
content_length=10,
|
||||
file_path=f"{suffix}.pdf",
|
||||
status=status,
|
||||
created_at="2024-01-01T00:00:00+00:00",
|
||||
updated_at="2024-01-01T00:00:00+00:00",
|
||||
metadata={},
|
||||
)
|
||||
|
||||
|
||||
class _FakeDocStatusStorage:
|
||||
def __init__(self):
|
||||
self.docs = {
|
||||
"processed-doc": _doc(DocStatus.PROCESSED, "processed"),
|
||||
"parsing-doc": _doc(DocStatus.PARSING, "parsing"),
|
||||
"analyzing-doc": _doc(DocStatus.ANALYZING, "analyzing"),
|
||||
}
|
||||
|
||||
async def get_docs_paginated(
|
||||
self,
|
||||
status_filter=None,
|
||||
status_filters=None,
|
||||
page=1,
|
||||
page_size=50,
|
||||
sort_field="updated_at",
|
||||
sort_direction="desc",
|
||||
):
|
||||
selected_statuses = DocStatusStorage.resolve_status_filter_values(
|
||||
status_filter=status_filter,
|
||||
status_filters=status_filters,
|
||||
)
|
||||
documents = [
|
||||
(doc_id, doc)
|
||||
for doc_id, doc in self.docs.items()
|
||||
if selected_statuses is None or doc.status.value in selected_statuses
|
||||
]
|
||||
return documents[:page_size], len(documents)
|
||||
|
||||
async def get_all_status_counts(self):
|
||||
return {"processed": 1, "parsing": 1, "analyzing": 1}
|
||||
|
||||
|
||||
_fake_doc_status = _FakeDocStatusStorage()
|
||||
_app = FastAPI()
|
||||
_app.include_router(
|
||||
create_document_routes(
|
||||
SimpleNamespace(doc_status=_fake_doc_status),
|
||||
SimpleNamespace(),
|
||||
api_key="test-key",
|
||||
)
|
||||
)
|
||||
_client = TestClient(_app)
|
||||
_headers = {"X-API-Key": "test-key"}
|
||||
|
||||
|
||||
def test_documents_paginated_accepts_status_filter():
|
||||
response = _client.post(
|
||||
"/documents/paginated",
|
||||
headers=_headers,
|
||||
json={
|
||||
"status_filter": "processed",
|
||||
"page": 1,
|
||||
"page_size": 10,
|
||||
"sort_field": "updated_at",
|
||||
"sort_direction": "desc",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
payload = response.json()
|
||||
assert payload["pagination"]["total_count"] == 1
|
||||
assert [doc["id"] for doc in payload["documents"]] == ["processed-doc"]
|
||||
|
||||
|
||||
def test_documents_paginated_status_filters_override_status_filter():
|
||||
response = _client.post(
|
||||
"/documents/paginated",
|
||||
headers=_headers,
|
||||
json={
|
||||
"status_filter": "processed",
|
||||
"status_filters": ["parsing", "analyzing"],
|
||||
"page": 1,
|
||||
"page_size": 10,
|
||||
"sort_field": "updated_at",
|
||||
"sort_direction": "desc",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
payload = response.json()
|
||||
assert payload["pagination"]["total_count"] == 2
|
||||
assert [doc["id"] for doc in payload["documents"]] == [
|
||||
"parsing-doc",
|
||||
"analyzing-doc",
|
||||
]
|
||||
@@ -0,0 +1,298 @@
|
||||
"""Pipeline-busy guard tests for graph mutation endpoints.
|
||||
|
||||
These tests verify that all 7 graph-mutation endpoints refuse to operate
|
||||
with HTTP 409 while the document pipeline is busy:
|
||||
|
||||
- POST /graph/entity/edit (graph_routes)
|
||||
- POST /graph/relation/edit (graph_routes)
|
||||
- POST /graph/entity/create (graph_routes)
|
||||
- POST /graph/relation/create (graph_routes)
|
||||
- POST /graph/entities/merge (graph_routes)
|
||||
- DELETE /graph/entity/delete (graph_routes)
|
||||
- DELETE /graph/relation/delete (graph_routes)
|
||||
|
||||
The guard logic itself lives in
|
||||
``lightrag.api.routers.document_routes.check_pipeline_busy_or_raise`` and is
|
||||
exercised both at the endpoint integration layer (via monkeypatch, no
|
||||
shared-storage dependency) and at the unit layer (against a real
|
||||
``pipeline_status`` namespace).
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
# Importing routers loads ``lightrag.api.config`` which parses ``sys.argv`` via
|
||||
# argparse. Stash argv so pytest's CLI flags don't trip the parser.
|
||||
_original_argv = sys.argv[:]
|
||||
sys.argv = [sys.argv[0]]
|
||||
_graph_routes = importlib.import_module("lightrag.api.routers.graph_routes")
|
||||
_document_routes = importlib.import_module("lightrag.api.routers.document_routes")
|
||||
sys.argv = _original_argv
|
||||
|
||||
create_graph_routes = _graph_routes.create_graph_routes
|
||||
create_document_routes = _document_routes.create_document_routes
|
||||
check_pipeline_busy_or_raise = _document_routes.check_pipeline_busy_or_raise
|
||||
|
||||
pytestmark = pytest.mark.offline
|
||||
|
||||
_API_KEY = "test-key"
|
||||
_HEADERS = {"X-API-Key": _API_KEY}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test scaffolding
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_mock_rag() -> SimpleNamespace:
|
||||
"""Build a minimal LightRAG stand-in with the 7 mutation methods stubbed.
|
||||
|
||||
Each ``AsyncMock`` returns a payload shaped enough to satisfy the
|
||||
endpoint's response model so the idle pass-through test can verify the
|
||||
full request path. Busy tests don't rely on these return values; the
|
||||
guard short-circuits before they're reached.
|
||||
"""
|
||||
return SimpleNamespace(
|
||||
workspace="",
|
||||
aedit_entity=AsyncMock(
|
||||
return_value={
|
||||
"entity_name": "Alice",
|
||||
"description": "updated",
|
||||
"operation_summary": {
|
||||
"merged": False,
|
||||
"merge_status": "not_attempted",
|
||||
"merge_error": None,
|
||||
"operation_status": "success",
|
||||
"target_entity": None,
|
||||
"final_entity": "Alice",
|
||||
"renamed": False,
|
||||
},
|
||||
}
|
||||
),
|
||||
aedit_relation=AsyncMock(return_value={"description": "updated"}),
|
||||
acreate_entity=AsyncMock(return_value={"entity_name": "Alice"}),
|
||||
acreate_relation=AsyncMock(return_value={"src_id": "a", "tgt_id": "b"}),
|
||||
amerge_entities=AsyncMock(return_value={"merged_entity": "Alice"}),
|
||||
adelete_by_entity=AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
status="success", message="deleted", doc_id="ignored"
|
||||
)
|
||||
),
|
||||
adelete_by_relation=AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
status="success", message="deleted", doc_id="ignored"
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _build_client(rag: SimpleNamespace) -> TestClient:
|
||||
app = FastAPI()
|
||||
app.include_router(create_graph_routes(rag, api_key=_API_KEY))
|
||||
app.include_router(create_document_routes(rag, SimpleNamespace(), api_key=_API_KEY))
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
async def _force_busy_guard(_rag) -> None:
|
||||
"""Stand-in for ``check_pipeline_busy_or_raise`` that always refuses."""
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=(
|
||||
"Pipeline is busy with another operation. "
|
||||
"Wait for the running job to finish before editing "
|
||||
"the knowledge graph."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _noop_guard(_rag) -> None:
|
||||
"""Stand-in for ``check_pipeline_busy_or_raise`` that always permits."""
|
||||
return None
|
||||
|
||||
|
||||
def _patch_guard(monkeypatch, replacement) -> None:
|
||||
"""Replace the guard reference in BOTH consumer modules.
|
||||
|
||||
``graph_routes`` re-binds the name via ``from .document_routes import ...``
|
||||
so patching only ``document_routes`` would miss the graph endpoints.
|
||||
"""
|
||||
monkeypatch.setattr(_graph_routes, "check_pipeline_busy_or_raise", replacement)
|
||||
monkeypatch.setattr(_document_routes, "check_pipeline_busy_or_raise", replacement)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Part A: endpoint integration -- guard refuses with 409
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_ENDPOINTS = [
|
||||
pytest.param(
|
||||
"POST",
|
||||
"/graph/entity/edit",
|
||||
{"entity_name": "Alice", "updated_data": {"description": "x"}},
|
||||
id="update_entity",
|
||||
),
|
||||
pytest.param(
|
||||
"POST",
|
||||
"/graph/relation/edit",
|
||||
{
|
||||
"source_id": "Alice",
|
||||
"target_id": "Bob",
|
||||
"updated_data": {"description": "x"},
|
||||
},
|
||||
id="update_relation",
|
||||
),
|
||||
pytest.param(
|
||||
"POST",
|
||||
"/graph/entity/create",
|
||||
{"entity_name": "Alice", "entity_data": {"description": "x"}},
|
||||
id="create_entity",
|
||||
),
|
||||
pytest.param(
|
||||
"POST",
|
||||
"/graph/relation/create",
|
||||
{
|
||||
"source_entity": "Alice",
|
||||
"target_entity": "Bob",
|
||||
"relation_data": {"description": "x"},
|
||||
},
|
||||
id="create_relation",
|
||||
),
|
||||
pytest.param(
|
||||
"POST",
|
||||
"/graph/entities/merge",
|
||||
{"entities_to_change": ["Alic"], "entity_to_change_into": "Alice"},
|
||||
id="merge_entities",
|
||||
),
|
||||
pytest.param(
|
||||
"DELETE",
|
||||
"/graph/entity/delete",
|
||||
{"entity_name": "Alice"},
|
||||
id="delete_entity",
|
||||
),
|
||||
pytest.param(
|
||||
"DELETE",
|
||||
"/graph/relation/delete",
|
||||
{"source_entity": "Alice", "target_entity": "Bob"},
|
||||
id="delete_relation",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method, path, body", _ENDPOINTS)
|
||||
def test_endpoint_refuses_with_409_when_pipeline_busy(method, path, body, monkeypatch):
|
||||
rag = _make_mock_rag()
|
||||
client = _build_client(rag)
|
||||
_patch_guard(monkeypatch, _force_busy_guard)
|
||||
|
||||
response = client.request(method, path, json=body, headers=_HEADERS)
|
||||
|
||||
assert response.status_code == 409, response.text
|
||||
payload = response.json()
|
||||
assert "Pipeline is busy" in payload["detail"]
|
||||
# Guard must short-circuit before the underlying mutation runs.
|
||||
for attr in (
|
||||
"aedit_entity",
|
||||
"aedit_relation",
|
||||
"acreate_entity",
|
||||
"acreate_relation",
|
||||
"amerge_entities",
|
||||
"adelete_by_entity",
|
||||
"adelete_by_relation",
|
||||
):
|
||||
getattr(rag, attr).assert_not_awaited()
|
||||
|
||||
|
||||
def test_endpoint_passes_through_when_pipeline_idle(monkeypatch):
|
||||
"""Sanity check: with an idle guard, the request reaches ``rag.aedit_entity``."""
|
||||
rag = _make_mock_rag()
|
||||
client = _build_client(rag)
|
||||
_patch_guard(monkeypatch, _noop_guard)
|
||||
|
||||
response = client.post(
|
||||
"/graph/entity/edit",
|
||||
json={"entity_name": "Alice", "updated_data": {"description": "x"}},
|
||||
headers=_HEADERS,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
rag.aedit_entity.assert_awaited_once()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Part B: helper unit -- against real pipeline_status namespace
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _with_pipeline_status(action):
|
||||
"""Bootstrap pipeline_status, run ``action(pipeline_status)``, then tear down.
|
||||
|
||||
``initialize_share_data`` is idempotent within a process but
|
||||
``finalize_share_data`` is required to release the Manager/lock state so
|
||||
repeated calls in subsequent tests start clean.
|
||||
"""
|
||||
from lightrag.kg.shared_storage import (
|
||||
finalize_share_data,
|
||||
get_namespace_data,
|
||||
initialize_pipeline_status,
|
||||
initialize_share_data,
|
||||
)
|
||||
|
||||
initialize_share_data()
|
||||
try:
|
||||
await initialize_pipeline_status(workspace="")
|
||||
pipeline_status = await get_namespace_data("pipeline_status", workspace="")
|
||||
await action(pipeline_status)
|
||||
finally:
|
||||
finalize_share_data()
|
||||
|
||||
|
||||
async def test_helper_raises_409_when_busy_flag_set():
|
||||
async def _do(pipeline_status):
|
||||
pipeline_status["busy"] = True
|
||||
rag = SimpleNamespace(workspace="")
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await check_pipeline_busy_or_raise(rag)
|
||||
assert exc_info.value.status_code == 409
|
||||
assert "Pipeline is busy" in exc_info.value.detail
|
||||
|
||||
await _with_pipeline_status(_do)
|
||||
|
||||
|
||||
async def test_helper_returns_silently_when_pipeline_idle():
|
||||
async def _do(pipeline_status):
|
||||
pipeline_status["busy"] = False
|
||||
rag = SimpleNamespace(workspace="")
|
||||
# Should not raise.
|
||||
await check_pipeline_busy_or_raise(rag)
|
||||
|
||||
await _with_pipeline_status(_do)
|
||||
|
||||
|
||||
async def test_helper_is_noop_when_pipeline_status_uninitialized():
|
||||
"""When pipeline_status namespace was never bootstrapped the helper must pass.
|
||||
|
||||
``get_namespace_data`` raises ``PipelineNotInitializedError`` when the
|
||||
pipeline_status namespace is missing (share data initialized but the
|
||||
pipeline namespace never created); the helper swallows that error so test
|
||||
rigs without an end-to-end RAG bootstrap stay green. Mirrors the existing
|
||||
contract of ``_acquire_destructive_busy``.
|
||||
"""
|
||||
from lightrag.kg.shared_storage import (
|
||||
finalize_share_data,
|
||||
initialize_share_data,
|
||||
)
|
||||
|
||||
initialize_share_data()
|
||||
try:
|
||||
rag = SimpleNamespace(workspace="__never_bootstrapped__")
|
||||
# Intentionally skip ``initialize_pipeline_status``: helper should
|
||||
# catch ``PipelineNotInitializedError`` and return silently.
|
||||
await check_pipeline_busy_or_raise(rag)
|
||||
finally:
|
||||
finalize_share_data()
|
||||
@@ -0,0 +1,174 @@
|
||||
"""
|
||||
Verify the /query and /query/stream endpoint response types.
|
||||
|
||||
Ensures:
|
||||
- /query → application/json (no streaming, backward-compatible)
|
||||
- /query/stream → application/x-ndjson
|
||||
"""
|
||||
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
|
||||
_ENV_VARS_TO_ISOLATE = (
|
||||
"LLM_BINDING",
|
||||
"EMBEDDING_BINDING",
|
||||
"LLM_BINDING_HOST",
|
||||
"LLM_BINDING_API_KEY",
|
||||
"LLM_MODEL",
|
||||
"EMBEDDING_BINDING_HOST",
|
||||
"EMBEDDING_BINDING_API_KEY",
|
||||
"EMBEDDING_MODEL",
|
||||
"LIGHTRAG_API_PREFIX",
|
||||
"LIGHTRAG_KV_STORAGE",
|
||||
"LIGHTRAG_VECTOR_STORAGE",
|
||||
"LIGHTRAG_GRAPH_STORAGE",
|
||||
"LIGHTRAG_DOC_STATUS_STORAGE",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolate_env(monkeypatch):
|
||||
for var in _ENV_VARS_TO_ISOLATE:
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
monkeypatch.setenv("LLM_BINDING", "ollama")
|
||||
monkeypatch.setenv("EMBEDDING_BINDING", "ollama")
|
||||
|
||||
|
||||
def _build_client():
|
||||
original_argv = sys.argv.copy()
|
||||
try:
|
||||
sys.argv = ["lightrag-server"]
|
||||
from lightrag.api.config import parse_args
|
||||
from lightrag.api.lightrag_server import create_app
|
||||
|
||||
args = parse_args()
|
||||
with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
|
||||
mock_rag.return_value = MagicMock()
|
||||
return TestClient(create_app(args))
|
||||
finally:
|
||||
sys.argv = original_argv
|
||||
|
||||
|
||||
class TestQueryRouteJsonOnly:
|
||||
"""The /query endpoint must stay JSON-only to preserve backward compatibility."""
|
||||
|
||||
def test_openapi_spec_declares_json_response(self):
|
||||
client = _build_client()
|
||||
response = client.get("/openapi.json")
|
||||
assert response.status_code == 200
|
||||
spec = response.json()
|
||||
|
||||
paths = spec.get("paths", {})
|
||||
query_path = paths.get("/query", {})
|
||||
assert query_path, "/query must be in OpenAPI paths"
|
||||
|
||||
post_op = query_path.get("post", {})
|
||||
responses = post_op.get("responses", {})
|
||||
ok_resp = responses.get("200", {})
|
||||
content = ok_resp.get("content", {})
|
||||
|
||||
# The /query endpoint must declare application/json — NOT ndjson
|
||||
assert "application/json" in content, (
|
||||
"/query must declare application/json in OpenAPI spec"
|
||||
)
|
||||
assert "application/x-ndjson" not in content, (
|
||||
"/query must NOT declare application/x-ndjson — streaming belongs to /query/stream"
|
||||
)
|
||||
|
||||
def test_query_route_exists_and_accepts_post(self):
|
||||
client = _build_client()
|
||||
# A minimal POST to /query should reach the route (it'll 422 or 500
|
||||
# since we don't have a real LLM, but it should NOT 404/405)
|
||||
response = client.post("/query", json={"query": "test", "mode": "mix"})
|
||||
assert response.status_code not in (
|
||||
404,
|
||||
405,
|
||||
), "/query route must exist and accept POST"
|
||||
|
||||
|
||||
class TestQueryStreamRoute:
|
||||
"""The /query/stream endpoint must serve application/x-ndjson."""
|
||||
|
||||
def test_openapi_spec_declares_ndjson_response(self):
|
||||
client = _build_client()
|
||||
response = client.get("/openapi.json")
|
||||
assert response.status_code == 200
|
||||
spec = response.json()
|
||||
|
||||
paths = spec.get("paths", {})
|
||||
stream_path = paths.get("/query/stream", {})
|
||||
assert stream_path, "/query/stream must be in OpenAPI paths"
|
||||
|
||||
post_op = stream_path.get("post", {})
|
||||
responses = post_op.get("responses", {})
|
||||
ok_resp = responses.get("200", {})
|
||||
content = ok_resp.get("content", {})
|
||||
|
||||
# The /query/stream endpoint must declare application/x-ndjson
|
||||
assert "application/x-ndjson" in content, (
|
||||
"/query/stream must declare application/x-ndjson in OpenAPI spec"
|
||||
)
|
||||
|
||||
def test_stream_route_exists_and_accepts_post(self):
|
||||
client = _build_client()
|
||||
response = client.post("/query/stream", json={"query": "test", "mode": "mix"})
|
||||
assert response.status_code not in (
|
||||
404,
|
||||
405,
|
||||
), "/query/stream route must exist and accept POST"
|
||||
|
||||
|
||||
class TestQueryStreamResponseContentType:
|
||||
"""When the mock LLM returns a non-streaming result, /query/stream must
|
||||
still set the correct Content-Type header."""
|
||||
|
||||
def test_stream_response_has_ndjson_content_type(self):
|
||||
"""Even without a real LLM, the streaming response must carry the
|
||||
correct media type header."""
|
||||
|
||||
original_argv = sys.argv.copy()
|
||||
try:
|
||||
sys.argv = ["lightrag-server"]
|
||||
from lightrag.api.config import parse_args
|
||||
from lightrag.api.lightrag_server import create_app
|
||||
|
||||
args = parse_args()
|
||||
|
||||
mock_rag = MagicMock()
|
||||
mock_result = {
|
||||
"llm_response": {
|
||||
"is_streaming": False,
|
||||
"content": "test response",
|
||||
},
|
||||
"data": {"references": []},
|
||||
}
|
||||
# Return a coroutine
|
||||
mock_rag.aquery_llm = MagicMock()
|
||||
|
||||
async def _fake_aquery(*a, **kw):
|
||||
return mock_result
|
||||
|
||||
mock_rag.aquery_llm.side_effect = _fake_aquery
|
||||
|
||||
with patch("lightrag.api.lightrag_server.LightRAG", return_value=mock_rag):
|
||||
app = create_app(args)
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post(
|
||||
"/query/stream",
|
||||
json={
|
||||
"query": "test",
|
||||
"mode": "mix",
|
||||
"include_references": True,
|
||||
},
|
||||
)
|
||||
content_type = response.headers.get("content-type", "")
|
||||
assert "application/x-ndjson" in content_type, (
|
||||
f"/query/stream must return application/x-ndjson, got: {content_type}"
|
||||
)
|
||||
finally:
|
||||
sys.argv = original_argv
|
||||
@@ -0,0 +1,48 @@
|
||||
"""Upload allowlist exposes the native markdown suffixes (.md / .textpack).
|
||||
|
||||
``DocumentManager.supported_extensions`` derives the uploadable suffix set
|
||||
live from the parser registry + routing: a suffix is advertised only when an
|
||||
unhinted ``x.<suffix>`` routes to an engine that actually supports it. These
|
||||
tests pin that ``.md`` and ``.textpack`` are uploadable (and routable), and
|
||||
that ``.textpack`` resolves to the native engine.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
# document_routes parses argv at import time; guard it like the sibling tests.
|
||||
_original_argv = sys.argv[:]
|
||||
sys.argv = [sys.argv[0]]
|
||||
_document_routes = importlib.import_module("lightrag.api.routers.document_routes")
|
||||
_parser_routing = importlib.import_module("lightrag.parser.routing")
|
||||
sys.argv = _original_argv
|
||||
|
||||
DocumentManager = _document_routes.DocumentManager
|
||||
resolve_file_parser_engine = _parser_routing.resolve_file_parser_engine
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def doc_manager(tmp_path, monkeypatch):
|
||||
# Clear any ambient LIGHTRAG_PARSER so the default routing applies:
|
||||
# .md -> legacy (opt-in for native), .textpack -> native.
|
||||
monkeypatch.delenv("LIGHTRAG_PARSER", raising=False)
|
||||
return DocumentManager(str(tmp_path))
|
||||
|
||||
|
||||
def test_textpack_and_md_are_uploadable(doc_manager):
|
||||
extensions = doc_manager.supported_extensions
|
||||
assert ".textpack" in extensions
|
||||
assert ".md" in extensions
|
||||
|
||||
|
||||
def test_bare_textpack_is_supported_file_via_native(doc_manager, monkeypatch):
|
||||
monkeypatch.delenv("LIGHTRAG_PARSER", raising=False)
|
||||
assert doc_manager.is_supported_file("note.textpack") is True
|
||||
assert resolve_file_parser_engine("note.textpack") == "native"
|
||||
|
||||
|
||||
def test_md_is_supported_file(doc_manager, monkeypatch):
|
||||
monkeypatch.delenv("LIGHTRAG_PARSER", raising=False)
|
||||
assert doc_manager.is_supported_file("doc.md") is True
|
||||
Reference in New Issue
Block a user