lix/mnote
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: init monorepo snapshot

Browse Source
This commit is contained in:
liaibo
2025-11-23 10:55:04 +08:00
commit c70ff52869
941 changed files with 246586 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
LightRAG/tests/test_aquery_data_endpoint.py
+768
View File
@@ -0,0 +1,768 @@
#!/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 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
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
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
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
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)
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)}")
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")
LightRAG/tests/test_curl_aquery_data.sh
+271
View File
@@ -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"
LightRAG/tests/test_graph_storage.py
+1301
View File
File diff suppressed because it is too large Load Diff
LightRAG/tests/test_lightrag_ollama_chat.py
+855
View File
@@ -0,0 +1,855 @@
"""
LightRAG Ollama Compatibility Interface Test Script
This script tests the LightRAG's Ollama compatibility interface, including:
1. Basic functionality tests (streaming and non-streaming responses)
2. Query mode tests (local, global, naive, hybrid)
3. Error handling tests (including streaming and non-streaming scenarios)
All responses use the JSON Lines format, complying with the Ollama API specification.
"""
import requests
import json
import argparse
import time
from typing import Dict, Any, Optional, List, Callable
from dataclasses import dataclass, asdict
from datetime import datetime
from pathlib import Path
from enum import Enum, auto
class ErrorCode(Enum):
"""Error codes for MCP errors"""
InvalidRequest = auto()
InternalError = auto()
class McpError(Exception):
"""Base exception class for MCP errors"""
def __init__(self, code: ErrorCode, message: str):
self.code = code
self.message = message
super().__init__(message)
DEFAULT_CONFIG = {
"server": {
"host": "localhost",
"port": 9621,
"model": "lightrag:latest",
"timeout": 300,
"max_retries": 1,
"retry_delay": 1,
},
"test_cases": {
"basic": {"query": "唐僧有几个徒弟"},
"generate": {"query": "电视剧西游记导演是谁"},
},
}
# Example conversation history for testing
EXAMPLE_CONVERSATION = [
{"role": "user", "content": "你好"},
{"role": "assistant", "content": "你好!我是一个AI助手,很高兴为你服务。"},
{"role": "user", "content": "Who are you?"},
{"role": "assistant", "content": "I'm a Knowledge base query assistant."},
]
class OutputControl:
"""Output control class, manages the verbosity of test output"""
_verbose: bool = False
@classmethod
def set_verbose(cls, verbose: bool) -> None:
cls._verbose = verbose
@classmethod
def is_verbose(cls) -> bool:
return cls._verbose
@dataclass
class TestResult:
"""Test result data class"""
name: str
success: bool
duration: float
error: Optional[str] = None
timestamp: str = ""
def __post_init__(self):
if not self.timestamp:
self.timestamp = datetime.now().isoformat()
class TestStats:
"""Test statistics"""
def __init__(self):
self.results: List[TestResult] = []
self.start_time = datetime.now()
def add_result(self, result: TestResult):
self.results.append(result)
def export_results(self, path: str = "test_results.json"):
"""Export test results to a JSON file
Args:
path: Output file path
"""
results_data = {
"start_time": self.start_time.isoformat(),
"end_time": datetime.now().isoformat(),
"results": [asdict(r) for r in self.results],
"summary": {
"total": len(self.results),
"passed": sum(1 for r in self.results if r.success),
"failed": sum(1 for r in self.results if not r.success),
"total_duration": sum(r.duration for r in self.results),
},
}
with open(path, "w", encoding="utf-8") as f:
json.dump(results_data, f, ensure_ascii=False, indent=2)
print(f"\nTest results saved to: {path}")
def print_summary(self):
total = len(self.results)
passed = sum(1 for r in self.results if r.success)
failed = total - passed
duration = sum(r.duration for r in self.results)
print("\n=== Test Summary ===")
print(f"Start time: {self.start_time.strftime('%Y-%m-%d %H:%M:%S')}")
print(f"Total duration: {duration:.2f} seconds")
print(f"Total tests: {total}")
print(f"Passed: {passed}")
print(f"Failed: {failed}")
if failed > 0:
print("\nFailed tests:")
for result in self.results:
if not result.success:
print(f"- {result.name}: {result.error}")
def make_request(
url: str, data: Dict[str, Any], stream: bool = False, check_status: bool = True
) -> requests.Response:
"""Send an HTTP request with retry mechanism
Args:
url: Request URL
data: Request data
stream: Whether to use streaming response
check_status: Whether to check HTTP status code (default: True)
Returns:
requests.Response: Response object
Raises:
requests.exceptions.RequestException: Request failed after all retries
requests.exceptions.HTTPError: HTTP status code is not 200 (when check_status is True)
"""
server_config = CONFIG["server"]
max_retries = server_config["max_retries"]
retry_delay = server_config["retry_delay"]
timeout = server_config["timeout"]
for attempt in range(max_retries):
try:
response = requests.post(url, json=data, stream=stream, timeout=timeout)
if check_status and response.status_code != 200:
response.raise_for_status()
return response
except requests.exceptions.RequestException as e:
if attempt == max_retries - 1: # Last retry
raise
print(f"\nRequest failed, retrying in {retry_delay} seconds: {str(e)}")
time.sleep(retry_delay)
def load_config() -> Dict[str, Any]:
"""Load configuration file
First try to load from config.json in the current directory,
if it doesn't exist, use the default configuration
Returns:
Configuration dictionary
"""
config_path = Path("config.json")
if config_path.exists():
with open(config_path, "r", encoding="utf-8") as f:
return json.load(f)
return DEFAULT_CONFIG
def print_json_response(data: Dict[str, Any], title: str = "", indent: int = 2) -> None:
"""Format and print JSON response data
Args:
data: Data dictionary to print
title: Title to print
indent: Number of spaces for JSON indentation
"""
if OutputControl.is_verbose():
if title:
print(f"\n=== {title} ===")
print(json.dumps(data, ensure_ascii=False, indent=indent))
# Global configuration
CONFIG = load_config()
def get_base_url(endpoint: str = "chat") -> str:
"""Return the base URL for specified endpoint
Args:
endpoint: API endpoint name (chat or generate)
Returns:
Complete URL for the endpoint
"""
server = CONFIG["server"]
return f"http://{server['host']}:{server['port']}/api/{endpoint}"
def create_chat_request_data(
content: str,
stream: bool = False,
model: str = None,
conversation_history: List[Dict[str, str]] = None,
) -> Dict[str, Any]:
"""Create chat request data
Args:
content: User message content
stream: Whether to use streaming response
model: Model name
conversation_history: List of previous conversation messages
history_turns: Number of history turns to include
Returns:
Dictionary containing complete chat request data
"""
messages = conversation_history or []
messages.append({"role": "user", "content": content})
return {
"model": model or CONFIG["server"]["model"],
"messages": messages,
"stream": stream,
}
def create_generate_request_data(
prompt: str,
system: str = None,
stream: bool = False,
model: str = None,
options: Dict[str, Any] = None,
) -> Dict[str, Any]:
"""Create generate request data
Args:
prompt: Generation prompt
system: System prompt
stream: Whether to use streaming response
model: Model name
options: Additional options
Returns:
Dictionary containing complete generate request data
"""
data = {
"model": model or CONFIG["server"]["model"],
"prompt": prompt,
"stream": stream,
}
if system:
data["system"] = system
if options:
data["options"] = options
return data
# Global test statistics
STATS = TestStats()
def run_test(func: Callable, name: str) -> None:
"""Run a test and record the results
Args:
func: Test function
name: Test name
"""
start_time = time.time()
try:
func()
duration = time.time() - start_time
STATS.add_result(TestResult(name, True, duration))
except Exception as e:
duration = time.time() - start_time
STATS.add_result(TestResult(name, False, duration, str(e)))
raise
def test_non_stream_chat() -> None:
"""Test non-streaming call to /api/chat endpoint"""
url = get_base_url()
# Send request with conversation history
data = create_chat_request_data(
CONFIG["test_cases"]["basic"]["query"],
stream=False,
conversation_history=EXAMPLE_CONVERSATION,
)
response = make_request(url, data)
# Print response
if OutputControl.is_verbose():
print("\n=== Non-streaming call response ===")
response_json = response.json()
# Print response content
print_json_response(
{"model": response_json["model"], "message": response_json["message"]},
"Response content",
)
def test_stream_chat() -> None:
"""Test streaming call to /api/chat endpoint
Use JSON Lines format to process streaming responses, each line is a complete JSON object.
Response format:
{
"model": "lightrag:latest",
"created_at": "2024-01-15T00:00:00Z",
"message": {
"role": "assistant",
"content": "Partial response content",
"images": null
},
"done": false
}
The last message will contain performance statistics, with done set to true.
"""
url = get_base_url()
# Send request with conversation history
data = create_chat_request_data(
CONFIG["test_cases"]["basic"]["query"],
stream=True,
conversation_history=EXAMPLE_CONVERSATION,
)
response = make_request(url, data, stream=True)
if OutputControl.is_verbose():
print("\n=== Streaming call response ===")
output_buffer = []
try:
for line in response.iter_lines():
if line: # Skip empty lines
try:
# Decode and parse JSON
data = json.loads(line.decode("utf-8"))
if data.get("done", True): # If it's the completion marker
if (
"total_duration" in data
): # Final performance statistics message
# print_json_response(data, "Performance statistics")
break
else: # Normal content message
message = data.get("message", {})
content = message.get("content", "")
if content: # Only collect non-empty content
output_buffer.append(content)
print(
content, end="", flush=True
) # Print content in real-time
except json.JSONDecodeError:
print("Error decoding JSON from response line")
finally:
response.close() # Ensure the response connection is closed
# Print a newline
print()
def test_query_modes() -> None:
"""Test different query mode prefixes
Supported query modes:
- /local: Local retrieval mode, searches only in highly relevant documents
- /global: Global retrieval mode, searches across all documents
- /naive: Naive mode, does not use any optimization strategies
- /hybrid: Hybrid mode (default), combines multiple strategies
- /mix: Mix mode
Each mode will return responses in the same format, but with different retrieval strategies.
"""
url = get_base_url()
modes = ["local", "global", "naive", "hybrid", "mix"]
for mode in modes:
if OutputControl.is_verbose():
print(f"\n=== Testing /{mode} mode ===")
data = create_chat_request_data(
f"/{mode} {CONFIG['test_cases']['basic']['query']}", stream=False
)
# Send request
response = make_request(url, data)
response_json = response.json()
# Print response content
print_json_response(
{"model": response_json["model"], "message": response_json["message"]}
)
def create_error_test_data(error_type: str) -> Dict[str, Any]:
"""Create request data for error testing
Args:
error_type: Error type, supported:
- empty_messages: Empty message list
- invalid_role: Invalid role field
- missing_content: Missing content field
Returns:
Request dictionary containing error data
"""
error_data = {
"empty_messages": {"model": "lightrag:latest", "messages": [], "stream": True},
"invalid_role": {
"model": "lightrag:latest",
"messages": [{"invalid_role": "user", "content": "Test message"}],
"stream": True,
},
"missing_content": {
"model": "lightrag:latest",
"messages": [{"role": "user"}],
"stream": True,
},
}
return error_data.get(error_type, error_data["empty_messages"])
def test_stream_error_handling() -> None:
"""Test error handling for streaming responses
Test scenarios:
1. Empty message list
2. Message format error (missing required fields)
Error responses should be returned immediately without establishing a streaming connection.
The status code should be 4xx, and detailed error information should be returned.
"""
url = get_base_url()
if OutputControl.is_verbose():
print("\n=== Testing streaming response error handling ===")
# Test empty message list
if OutputControl.is_verbose():
print("\n--- Testing empty message list (streaming) ---")
data = create_error_test_data("empty_messages")
response = make_request(url, data, stream=True, check_status=False)
print(f"Status code: {response.status_code}")
if response.status_code != 200:
print_json_response(response.json(), "Error message")
response.close()
# Test invalid role field
if OutputControl.is_verbose():
print("\n--- Testing invalid role field (streaming) ---")
data = create_error_test_data("invalid_role")
response = make_request(url, data, stream=True, check_status=False)
print(f"Status code: {response.status_code}")
if response.status_code != 200:
print_json_response(response.json(), "Error message")
response.close()
# Test missing content field
if OutputControl.is_verbose():
print("\n--- Testing missing content field (streaming) ---")
data = create_error_test_data("missing_content")
response = make_request(url, data, stream=True, check_status=False)
print(f"Status code: {response.status_code}")
if response.status_code != 200:
print_json_response(response.json(), "Error message")
response.close()
def test_error_handling() -> None:
"""Test error handling for non-streaming responses
Test scenarios:
1. Empty message list
2. Message format error (missing required fields)
Error response format:
{
"detail": "Error description"
}
All errors should return appropriate HTTP status codes and clear error messages.
"""
url = get_base_url()
if OutputControl.is_verbose():
print("\n=== Testing error handling ===")
# Test empty message list
if OutputControl.is_verbose():
print("\n--- Testing empty message list ---")
data = create_error_test_data("empty_messages")
data["stream"] = False # Change to non-streaming mode
response = make_request(url, data, check_status=False)
print(f"Status code: {response.status_code}")
print_json_response(response.json(), "Error message")
# Test invalid role field
if OutputControl.is_verbose():
print("\n--- Testing invalid role field ---")
data = create_error_test_data("invalid_role")
data["stream"] = False # Change to non-streaming mode
response = make_request(url, data, check_status=False)
print(f"Status code: {response.status_code}")
print_json_response(response.json(), "Error message")
# Test missing content field
if OutputControl.is_verbose():
print("\n--- Testing missing content field ---")
data = create_error_test_data("missing_content")
data["stream"] = False # Change to non-streaming mode
response = make_request(url, data, check_status=False)
print(f"Status code: {response.status_code}")
print_json_response(response.json(), "Error message")
def test_non_stream_generate() -> None:
"""Test non-streaming call to /api/generate endpoint"""
url = get_base_url("generate")
data = create_generate_request_data(
CONFIG["test_cases"]["generate"]["query"], stream=False
)
# Send request
response = make_request(url, data)
# Print response
if OutputControl.is_verbose():
print("\n=== Non-streaming generate response ===")
response_json = response.json()
# Print response content
print(json.dumps(response_json, ensure_ascii=False, indent=2))
def test_stream_generate() -> None:
"""Test streaming call to /api/generate endpoint"""
url = get_base_url("generate")
data = create_generate_request_data(
CONFIG["test_cases"]["generate"]["query"], stream=True
)
# Send request and get streaming response
response = make_request(url, data, stream=True)
if OutputControl.is_verbose():
print("\n=== Streaming generate response ===")
output_buffer = []
try:
for line in response.iter_lines():
if line: # Skip empty lines
try:
# Decode and parse JSON
data = json.loads(line.decode("utf-8"))
if data.get("done", True): # If it's the completion marker
if (
"total_duration" in data
): # Final performance statistics message
break
else: # Normal content message
content = data.get("response", "")
if content: # Only collect non-empty content
output_buffer.append(content)
print(
content, end="", flush=True
) # Print content in real-time
except json.JSONDecodeError:
print("Error decoding JSON from response line")
finally:
response.close() # Ensure the response connection is closed
# Print a newline
print()
def test_generate_with_system() -> None:
"""Test generate with system prompt"""
url = get_base_url("generate")
data = create_generate_request_data(
CONFIG["test_cases"]["generate"]["query"],
system="你是一个知识渊博的助手",
stream=False,
)
# Send request
response = make_request(url, data)
# Print response
if OutputControl.is_verbose():
print("\n=== Generate with system prompt response ===")
response_json = response.json()
# Print response content
print_json_response(
{
"model": response_json["model"],
"response": response_json["response"],
"done": response_json["done"],
},
"Response content",
)
def test_generate_error_handling() -> None:
"""Test error handling for generate endpoint"""
url = get_base_url("generate")
# Test empty prompt
if OutputControl.is_verbose():
print("\n=== Testing empty prompt ===")
data = create_generate_request_data("", stream=False)
response = make_request(url, data, check_status=False)
print(f"Status code: {response.status_code}")
print_json_response(response.json(), "Error message")
# Test invalid options
if OutputControl.is_verbose():
print("\n=== Testing invalid options ===")
data = create_generate_request_data(
CONFIG["test_cases"]["basic"]["query"],
options={"invalid_option": "value"},
stream=False,
)
response = make_request(url, data, check_status=False)
print(f"Status code: {response.status_code}")
print_json_response(response.json(), "Error message")
def test_generate_concurrent() -> None:
"""Test concurrent generate requests"""
import asyncio
import aiohttp
from contextlib import asynccontextmanager
@asynccontextmanager
async def get_session():
async with aiohttp.ClientSession() as session:
yield session
async def make_request(session, prompt: str, request_id: int):
url = get_base_url("generate")
data = create_generate_request_data(prompt, stream=False)
try:
async with session.post(url, json=data) as response:
if response.status != 200:
error_msg = (
f"Request {request_id} failed with status {response.status}"
)
if OutputControl.is_verbose():
print(f"\n{error_msg}")
raise McpError(ErrorCode.InternalError, error_msg)
result = await response.json()
if "error" in result:
error_msg = (
f"Request {request_id} returned error: {result['error']}"
)
if OutputControl.is_verbose():
print(f"\n{error_msg}")
raise McpError(ErrorCode.InternalError, error_msg)
return result
except Exception as e:
error_msg = f"Request {request_id} failed: {str(e)}"
if OutputControl.is_verbose():
print(f"\n{error_msg}")
raise McpError(ErrorCode.InternalError, error_msg)
async def run_concurrent_requests():
prompts = ["第一个问题", "第二个问题", "第三个问题", "第四个问题", "第五个问题"]
async with get_session() as session:
tasks = [
make_request(session, prompt, i + 1) for i, prompt in enumerate(prompts)
]
results = await asyncio.gather(*tasks, return_exceptions=True)
success_results = []
error_messages = []
for i, result in enumerate(results):
if isinstance(result, Exception):
error_messages.append(f"Request {i+1} failed: {str(result)}")
else:
success_results.append((i + 1, result))
if error_messages:
for req_id, result in success_results:
if OutputControl.is_verbose():
print(f"\nRequest {req_id} succeeded:")
print_json_response(result)
error_summary = "\n".join(error_messages)
raise McpError(
ErrorCode.InternalError,
f"Some concurrent requests failed:\n{error_summary}",
)
return results
if OutputControl.is_verbose():
print("\n=== Testing concurrent generate requests ===")
# Run concurrent requests
try:
results = asyncio.run(run_concurrent_requests())
# all success, print out results
for i, result in enumerate(results, 1):
print(f"\nRequest {i} result:")
print_json_response(result)
except McpError:
# error message already printed
raise
def get_test_cases() -> Dict[str, Callable]:
"""Get all available test cases
Returns:
A dictionary mapping test names to test functions
"""
return {
"non_stream": test_non_stream_chat,
"stream": test_stream_chat,
"modes": test_query_modes,
"errors": test_error_handling,
"stream_errors": test_stream_error_handling,
"non_stream_generate": test_non_stream_generate,
"stream_generate": test_stream_generate,
"generate_with_system": test_generate_with_system,
"generate_errors": test_generate_error_handling,
"generate_concurrent": test_generate_concurrent,
}
def create_default_config():
"""Create a default configuration file"""
config_path = Path("config.json")
if not config_path.exists():
with open(config_path, "w", encoding="utf-8") as f:
json.dump(DEFAULT_CONFIG, f, ensure_ascii=False, indent=2)
print(f"Default configuration file created: {config_path}")
def parse_args() -> argparse.Namespace:
"""Parse command line arguments"""
parser = argparse.ArgumentParser(
description="LightRAG Ollama Compatibility Interface Testing",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Configuration file (config.json):
{
"server": {
"host": "localhost", # Server address
"port": 9621, # Server port
"model": "lightrag:latest" # Default model name
},
"test_cases": {
"basic": {
"query": "Test query", # Basic query text
"stream_query": "Stream query" # Stream query text
}
}
}
""",
)
parser.add_argument(
"-q",
"--quiet",
action="store_true",
help="Silent mode, only display test result summary",
)
parser.add_argument(
"-a",
"--ask",
type=str,
help="Specify query content, which will override the query settings in the configuration file",
)
parser.add_argument(
"--init-config", action="store_true", help="Create default configuration file"
)
parser.add_argument(
"--output",
type=str,
default="",
help="Test result output file path, default is not to output to a file",
)
parser.add_argument(
"--tests",
nargs="+",
choices=list(get_test_cases().keys()) + ["all"],
default=["all"],
help="Test cases to run, options: %(choices)s. Use 'all' to run all tests （except error tests)",
)
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
# Set output mode
OutputControl.set_verbose(not args.quiet)
# If query content is specified, update the configuration
if args.ask:
CONFIG["test_cases"]["basic"]["query"] = args.ask
# If specified to create a configuration file
if args.init_config:
create_default_config()
exit(0)
test_cases = get_test_cases()
try:
if "all" in args.tests:
# Run all tests except error handling tests
if OutputControl.is_verbose():
print("\n【Chat API Tests】")
run_test(test_non_stream_chat, "Non-streaming Chat Test")
run_test(test_stream_chat, "Streaming Chat Test")
run_test(test_query_modes, "Chat Query Mode Test")
if OutputControl.is_verbose():
print("\n【Generate API Tests】")
run_test(test_non_stream_generate, "Non-streaming Generate Test")
run_test(test_stream_generate, "Streaming Generate Test")
run_test(test_generate_with_system, "Generate with System Prompt Test")
run_test(test_generate_concurrent, "Generate Concurrent Test")
else:
# Run specified tests
for test_name in args.tests:
if OutputControl.is_verbose():
print(f"\n【Running Test: {test_name}】")
run_test(test_cases[test_name], test_name)
except Exception as e:
print(f"\nAn error occurred: {str(e)}")
finally:
# Print test statistics
STATS.print_summary()
# If an output file path is specified, export the results
if args.output:
STATS.export_results(args.output)
LightRAG/tests/test_postgres_retry_integration.py
+338
View File
@@ -0,0 +1,338 @@
"""
Integration test suite for PostgreSQL retry mechanism using real database.
This test suite connects to a real PostgreSQL database using credentials from .env
and tests the retry mechanism with actual network failures.
Prerequisites:
1. PostgreSQL server running and accessible
2. .env file with POSTGRES_* configuration
3. asyncpg installed: pip install asyncpg
"""
import pytest
import asyncio
import os
import time
from dotenv import load_dotenv
from unittest.mock import patch
from lightrag.kg.postgres_impl import PostgreSQLDB
asyncpg = pytest.importorskip("asyncpg")
# Load environment variables
load_dotenv(dotenv_path=".env", override=False)
class TestPostgresRetryIntegration:
"""Integration tests for PostgreSQL retry mechanism with real database."""
@pytest.fixture
def db_config(self):
"""Load database configuration from environment variables."""
return {
"host": os.getenv("POSTGRES_HOST", "localhost"),
"port": int(os.getenv("POSTGRES_PORT", "5432")),
"user": os.getenv("POSTGRES_USER", "postgres"),
"password": os.getenv("POSTGRES_PASSWORD", ""),
"database": os.getenv("POSTGRES_DATABASE", "postgres"),
"workspace": os.getenv("POSTGRES_WORKSPACE", "test_retry"),
"max_connections": int(os.getenv("POSTGRES_MAX_CONNECTIONS", "10")),
# Connection retry configuration
"connection_retry_attempts": min(
10, int(os.getenv("POSTGRES_CONNECTION_RETRIES", "3"))
),
"connection_retry_backoff": min(
5.0, float(os.getenv("POSTGRES_CONNECTION_RETRY_BACKOFF", "0.5"))
),
"connection_retry_backoff_max": min(
60.0, float(os.getenv("POSTGRES_CONNECTION_RETRY_BACKOFF_MAX", "5.0"))
),
"pool_close_timeout": min(
30.0, float(os.getenv("POSTGRES_POOL_CLOSE_TIMEOUT", "5.0"))
),
}
@pytest.fixture
def test_env(self, monkeypatch):
"""Set up test environment variables for retry configuration."""
monkeypatch.setenv("POSTGRES_CONNECTION_RETRIES", "3")
monkeypatch.setenv("POSTGRES_CONNECTION_RETRY_BACKOFF", "0.5")
monkeypatch.setenv("POSTGRES_CONNECTION_RETRY_BACKOFF_MAX", "2.0")
monkeypatch.setenv("POSTGRES_POOL_CLOSE_TIMEOUT", "3.0")
@pytest.mark.asyncio
async def test_real_connection_success(self, db_config, test_env):
"""
Test successful connection to real PostgreSQL database.
This validates that:
1. Database credentials are correct
2. Connection pool initializes properly
3. Basic query works
"""
print("\n" + "=" * 80)
print("INTEGRATION TEST 1: Real Database Connection")
print("=" * 80)
print(
f" → Connecting to {db_config['host']}:{db_config['port']}/{db_config['database']}"
)
db = PostgreSQLDB(db_config)
try:
# Initialize database connection
await db.initdb()
print(" ✓ Connection successful")
# Test simple query
result = await db.query("SELECT 1 as test", multirows=False)
assert result is not None
assert result.get("test") == 1
print(" ✓ Query executed successfully")
print("\n✅ Test passed: Real database connection works")
print("=" * 80)
finally:
if db.pool:
await db.pool.close()
@pytest.mark.asyncio
async def test_simulated_transient_error_with_real_db(self, db_config, test_env):
"""
Test retry mechanism with simulated transient errors on real database.
Simulates connection failures on first 2 attempts, then succeeds.
"""
print("\n" + "=" * 80)
print("INTEGRATION TEST 2: Simulated Transient Errors")
print("=" * 80)
db = PostgreSQLDB(db_config)
attempt_count = {"value": 0}
# Original create_pool function
original_create_pool = asyncpg.create_pool
async def mock_create_pool_with_failures(*args, **kwargs):
"""Mock that fails first 2 times, then calls real create_pool."""
attempt_count["value"] += 1
print(f" → Connection attempt {attempt_count['value']}")
if attempt_count["value"] <= 2:
print(" ✗ Simulating connection failure")
raise asyncpg.exceptions.ConnectionFailureError(
f"Simulated failure on attempt {attempt_count['value']}"
)
print(" ✓ Allowing real connection")
return await original_create_pool(*args, **kwargs)
try:
# Patch create_pool to simulate failures
with patch(
"asyncpg.create_pool", side_effect=mock_create_pool_with_failures
):
await db.initdb()
assert (
attempt_count["value"] == 3
), f"Expected 3 attempts, got {attempt_count['value']}"
assert db.pool is not None, "Pool should be initialized after retries"
# Verify database is actually working
result = await db.query("SELECT 1 as test", multirows=False)
assert result.get("test") == 1
print(
f"\n✅ Test passed: Retry mechanism worked, connected after {attempt_count['value']} attempts"
)
print("=" * 80)
finally:
if db.pool:
await db.pool.close()
@pytest.mark.asyncio
async def test_query_retry_with_real_db(self, db_config, test_env):
"""
Test query-level retry with simulated connection issues.
Tests that queries retry on transient failures by simulating
a temporary database unavailability.
"""
print("\n" + "=" * 80)
print("INTEGRATION TEST 3: Query-Level Retry")
print("=" * 80)
db = PostgreSQLDB(db_config)
try:
# First initialize normally
await db.initdb()
print(" ✓ Database initialized")
# Close the pool to simulate connection loss
print(" → Simulating connection loss (closing pool)...")
await db.pool.close()
db.pool = None
# Now query should trigger pool recreation and retry
print(" → Attempting query (should auto-reconnect)...")
result = await db.query("SELECT 1 as test", multirows=False)
assert result.get("test") == 1, "Query should succeed after reconnection"
assert db.pool is not None, "Pool should be recreated"
print(" ✓ Query succeeded after automatic reconnection")
print("\n✅ Test passed: Auto-reconnection works correctly")
print("=" * 80)
finally:
if db.pool:
await db.pool.close()
@pytest.mark.asyncio
async def test_concurrent_queries_with_real_db(self, db_config, test_env):
"""
Test concurrent queries to validate thread safety and connection pooling.
Runs multiple concurrent queries to ensure no deadlocks or race conditions.
"""
print("\n" + "=" * 80)
print("INTEGRATION TEST 4: Concurrent Queries")
print("=" * 80)
db = PostgreSQLDB(db_config)
try:
await db.initdb()
print(" ✓ Database initialized")
# Launch 10 concurrent queries
num_queries = 10
print(f" → Launching {num_queries} concurrent queries...")
async def run_query(query_id):
result = await db.query(
f"SELECT {query_id} as id, pg_sleep(0.1)", multirows=False
)
return result.get("id")
start_time = time.time()
tasks = [run_query(i) for i in range(num_queries)]
results = await asyncio.gather(*tasks, return_exceptions=True)
elapsed = time.time() - start_time
# Check results
successful = sum(1 for r in results if not isinstance(r, Exception))
failed = sum(1 for r in results if isinstance(r, Exception))
print(f" → Completed in {elapsed:.2f}s")
print(f" → Results: {successful} successful, {failed} failed")
assert (
successful == num_queries
), f"All {num_queries} queries should succeed"
assert failed == 0, "No queries should fail"
print("\n✅ Test passed: All concurrent queries succeeded, no deadlocks")
print("=" * 80)
finally:
if db.pool:
await db.pool.close()
@pytest.mark.asyncio
async def test_pool_close_timeout_real(self, db_config, test_env):
"""
Test pool close timeout protection with real database.
"""
print("\n" + "=" * 80)
print("INTEGRATION TEST 5: Pool Close Timeout")
print("=" * 80)
db = PostgreSQLDB(db_config)
try:
await db.initdb()
print(" ✓ Database initialized")
# Trigger pool reset (which includes close)
print(" → Triggering pool reset...")
start_time = time.time()
await db._reset_pool()
elapsed = time.time() - start_time
print(f" ✓ Pool reset completed in {elapsed:.2f}s")
assert db.pool is None, "Pool should be None after reset"
assert (
elapsed < db.pool_close_timeout + 1
), "Reset should complete within timeout"
print("\n✅ Test passed: Pool reset handled correctly")
print("=" * 80)
finally:
# Already closed in test
pass
@pytest.mark.asyncio
async def test_configuration_from_env(self, db_config):
"""
Test that configuration is correctly loaded from environment variables.
"""
print("\n" + "=" * 80)
print("INTEGRATION TEST 6: Environment Configuration")
print("=" * 80)
db = PostgreSQLDB(db_config)
print(" → Configuration loaded:")
print(f" • Host: {db.host}")
print(f" • Port: {db.port}")
print(f" • Database: {db.database}")
print(f" • User: {db.user}")
print(f" • Workspace: {db.workspace}")
print(f" • Max Connections: {db.max}")
print(f" • Retry Attempts: {db.connection_retry_attempts}")
print(f" • Retry Backoff: {db.connection_retry_backoff}s")
print(f" • Max Backoff: {db.connection_retry_backoff_max}s")
print(f" • Pool Close Timeout: {db.pool_close_timeout}s")
# Verify required fields are present
assert db.host, "Host should be configured"
assert db.port, "Port should be configured"
assert db.user, "User should be configured"
assert db.database, "Database should be configured"
print("\n✅ Test passed: All configuration loaded correctly from .env")
print("=" * 80)
def run_integration_tests():
"""Run all integration tests with detailed output."""
print("\n" + "=" * 80)
print("POSTGRESQL RETRY MECHANISM - INTEGRATION TESTS")
print("Testing with REAL database from .env configuration")
print("=" * 80)
# Check if database configuration exists
if not os.getenv("POSTGRES_HOST"):
print("\n⚠️ WARNING: No POSTGRES_HOST in .env file")
print("Please ensure .env file exists with PostgreSQL configuration.")
return
print("\nRunning integration tests...\n")
# Run pytest with verbose output
pytest.main(
[
__file__,
"-v",
"-s", # Don't capture output
"--tb=short", # Short traceback format
"--color=yes",
"-x", # Stop on first failure
]
)
if __name__ == "__main__":
run_integration_tests()
LightRAG/tests/test_write_json_optimization.py
+387
View File
@@ -0,0 +1,387 @@
"""
Test suite for write_json optimization
This test verifies:
1. Fast path works for clean data (no sanitization)
2. Slow path applies sanitization for dirty data
3. Sanitization is done during encoding (memory-efficient)
4. Reloading updates shared memory with cleaned data
"""
import os
import json
import tempfile
from lightrag.utils import write_json, load_json, SanitizingJSONEncoder
class TestWriteJsonOptimization:
"""Test write_json optimization with two-stage approach"""
def test_fast_path_clean_data(self):
"""Test that clean data takes the fast path without sanitization"""
clean_data = {
"name": "John Doe",
"age": 30,
"items": ["apple", "banana", "cherry"],
"nested": {"key": "value", "number": 42},
}
with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
temp_file = f.name
try:
# Write clean data - should return False (no sanitization)
needs_reload = write_json(clean_data, temp_file)
assert not needs_reload, "Clean data should not require sanitization"
# Verify data was written correctly
loaded_data = load_json(temp_file)
assert loaded_data == clean_data, "Loaded data should match original"
finally:
os.unlink(temp_file)
def test_slow_path_dirty_data(self):
"""Test that dirty data triggers sanitization"""
# Create data with surrogate characters (U+D800 to U+DFFF)
dirty_string = "Hello\ud800World" # Contains surrogate character
dirty_data = {"text": dirty_string, "number": 123}
with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
temp_file = f.name
try:
# Write dirty data - should return True (sanitization applied)
needs_reload = write_json(dirty_data, temp_file)
assert needs_reload, "Dirty data should trigger sanitization"
# Verify data was written and sanitized
loaded_data = load_json(temp_file)
assert loaded_data is not None, "Data should be written"
assert loaded_data["number"] == 123, "Clean fields should remain unchanged"
# Surrogate character should be removed
assert (
"\ud800" not in loaded_data["text"]
), "Surrogate character should be removed"
finally:
os.unlink(temp_file)
def test_sanitizing_encoder_removes_surrogates(self):
"""Test that SanitizingJSONEncoder removes surrogate characters"""
data_with_surrogates = {
"text": "Hello\ud800\udc00World", # Contains surrogate pair
"clean": "Clean text",
"nested": {"dirty_key\ud801": "value", "clean_key": "clean\ud802value"},
}
# Encode using custom encoder
encoded = json.dumps(
data_with_surrogates, cls=SanitizingJSONEncoder, ensure_ascii=False
)
# Verify no surrogate characters in output
assert "\ud800" not in encoded, "Surrogate U+D800 should be removed"
assert "\udc00" not in encoded, "Surrogate U+DC00 should be removed"
assert "\ud801" not in encoded, "Surrogate U+D801 should be removed"
assert "\ud802" not in encoded, "Surrogate U+D802 should be removed"
# Verify clean parts remain
assert "Clean text" in encoded, "Clean text should remain"
assert "clean_key" in encoded, "Clean keys should remain"
def test_nested_structure_sanitization(self):
"""Test sanitization of deeply nested structures"""
nested_data = {
"level1": {
"level2": {
"level3": {"dirty": "text\ud800here", "clean": "normal text"},
"list": ["item1", "item\ud801dirty", "item3"],
}
}
}
with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
temp_file = f.name
try:
needs_reload = write_json(nested_data, temp_file)
assert needs_reload, "Nested dirty data should trigger sanitization"
# Verify nested structure is preserved
loaded_data = load_json(temp_file)
assert "level1" in loaded_data
assert "level2" in loaded_data["level1"]
assert "level3" in loaded_data["level1"]["level2"]
# Verify surrogates are removed
dirty_text = loaded_data["level1"]["level2"]["level3"]["dirty"]
assert "\ud800" not in dirty_text, "Nested surrogate should be removed"
# Verify list items are sanitized
list_items = loaded_data["level1"]["level2"]["list"]
assert (
"\ud801" not in list_items[1]
), "List item surrogates should be removed"
finally:
os.unlink(temp_file)
def test_unicode_non_characters_removed(self):
"""Test that Unicode non-characters (U+FFFE, U+FFFF) don't cause encoding errors
Note: U+FFFE and U+FFFF are valid UTF-8 characters (though discouraged),
so they don't trigger sanitization. They only get removed when explicitly
using the SanitizingJSONEncoder.
"""
data_with_nonchars = {"text1": "Hello\ufffeWorld", "text2": "Test\uffffString"}
with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
temp_file = f.name
try:
# These characters are valid UTF-8, so they take the fast path
needs_reload = write_json(data_with_nonchars, temp_file)
assert not needs_reload, "U+FFFE/U+FFFF are valid UTF-8 characters"
loaded_data = load_json(temp_file)
# They're written as-is in the fast path
assert loaded_data == data_with_nonchars
finally:
os.unlink(temp_file)
def test_mixed_clean_dirty_data(self):
"""Test data with both clean and dirty fields"""
mixed_data = {
"clean_field": "This is perfectly fine",
"dirty_field": "This has\ud800issues",
"number": 42,
"boolean": True,
"null_value": None,
"clean_list": [1, 2, 3],
"dirty_list": ["clean", "dirty\ud801item"],
}
with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
temp_file = f.name
try:
needs_reload = write_json(mixed_data, temp_file)
assert (
needs_reload
), "Mixed data with dirty fields should trigger sanitization"
loaded_data = load_json(temp_file)
# Clean fields should remain unchanged
assert loaded_data["clean_field"] == "This is perfectly fine"
assert loaded_data["number"] == 42
assert loaded_data["boolean"]
assert loaded_data["null_value"] is None
assert loaded_data["clean_list"] == [1, 2, 3]
# Dirty fields should be sanitized
assert "\ud800" not in loaded_data["dirty_field"]
assert "\ud801" not in loaded_data["dirty_list"][1]
finally:
os.unlink(temp_file)
def test_empty_and_none_strings(self):
"""Test handling of empty and None values"""
data = {
"empty": "",
"none": None,
"zero": 0,
"false": False,
"empty_list": [],
"empty_dict": {},
}
with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
temp_file = f.name
try:
needs_reload = write_json(data, temp_file)
assert (
not needs_reload
), "Clean empty values should not trigger sanitization"
loaded_data = load_json(temp_file)
assert loaded_data == data, "Empty/None values should be preserved"
finally:
os.unlink(temp_file)
def test_specific_surrogate_udc9a(self):
"""Test specific surrogate character \\udc9a mentioned in the issue"""
# Test the exact surrogate character from the error message:
# UnicodeEncodeError: 'utf-8' codec can't encode character '\\udc9a'
data_with_udc9a = {
"text": "Some text with surrogate\udc9acharacter",
"position": 201, # As mentioned in the error
"clean_field": "Normal text",
}
with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
temp_file = f.name
try:
# Write data - should trigger sanitization
needs_reload = write_json(data_with_udc9a, temp_file)
assert needs_reload, "Data with \\udc9a should trigger sanitization"
# Verify surrogate was removed
loaded_data = load_json(temp_file)
assert loaded_data is not None
assert "\udc9a" not in loaded_data["text"], "\\udc9a should be removed"
assert (
loaded_data["clean_field"] == "Normal text"
), "Clean fields should remain"
finally:
os.unlink(temp_file)
def test_migration_with_surrogate_sanitization(self):
"""Test that migration process handles surrogate characters correctly
This test simulates the scenario where legacy cache contains surrogate
characters and ensures they are cleaned during migration.
"""
# Simulate legacy cache data with surrogate characters
legacy_data_with_surrogates = {
"cache_entry_1": {
"return": "Result with\ud800surrogate",
"cache_type": "extract",
"original_prompt": "Some\udc9aprompt",
},
"cache_entry_2": {
"return": "Clean result",
"cache_type": "query",
"original_prompt": "Clean prompt",
},
}
with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
temp_file = f.name
try:
# First write the dirty data directly (simulating legacy cache file)
# Use custom encoder to force write even with surrogates
with open(temp_file, "w", encoding="utf-8") as f:
json.dump(
legacy_data_with_surrogates,
f,
cls=SanitizingJSONEncoder,
ensure_ascii=False,
)
# Load and verify surrogates were cleaned during initial write
loaded_data = load_json(temp_file)
assert loaded_data is not None
# The data should be sanitized
assert (
"\ud800" not in loaded_data["cache_entry_1"]["return"]
), "Surrogate in return should be removed"
assert (
"\udc9a" not in loaded_data["cache_entry_1"]["original_prompt"]
), "Surrogate in prompt should be removed"
# Clean data should remain unchanged
assert (
loaded_data["cache_entry_2"]["return"] == "Clean result"
), "Clean data should remain"
finally:
os.unlink(temp_file)
def test_empty_values_after_sanitization(self):
"""Test that data with empty values after sanitization is properly handled
Critical edge case: When sanitization results in data with empty string values,
we must use 'if cleaned_data is not None' instead of 'if cleaned_data' to ensure
proper reload, since truthy check on dict depends on content, not just existence.
"""
# Create data where ALL values are only surrogate characters
all_dirty_data = {
"key1": "\ud800\udc00\ud801",
"key2": "\ud802\ud803",
}
with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
temp_file = f.name
try:
# Write dirty data - should trigger sanitization
needs_reload = write_json(all_dirty_data, temp_file)
assert needs_reload, "All-dirty data should trigger sanitization"
# Load the sanitized data
cleaned_data = load_json(temp_file)
# Critical assertions for the edge case
assert cleaned_data is not None, "Cleaned data should not be None"
# Sanitization removes surrogates but preserves keys with empty values
assert cleaned_data == {
"key1": "",
"key2": "",
}, "Surrogates should be removed, keys preserved"
# This dict is truthy because it has keys (even with empty values)
assert cleaned_data, "Dict with keys is truthy"
# Test the actual edge case: empty dict
empty_data = {}
needs_reload2 = write_json(empty_data, temp_file)
assert not needs_reload2, "Empty dict is clean"
reloaded_empty = load_json(temp_file)
assert reloaded_empty is not None, "Empty dict should not be None"
assert reloaded_empty == {}, "Empty dict should remain empty"
assert (
not reloaded_empty
), "Empty dict evaluates to False (the critical check)"
finally:
os.unlink(temp_file)
if __name__ == "__main__":
# Run tests
test = TestWriteJsonOptimization()
print("Running test_fast_path_clean_data...")
test.test_fast_path_clean_data()
print("✓ Passed")
print("Running test_slow_path_dirty_data...")
test.test_slow_path_dirty_data()
print("✓ Passed")
print("Running test_sanitizing_encoder_removes_surrogates...")
test.test_sanitizing_encoder_removes_surrogates()
print("✓ Passed")
print("Running test_nested_structure_sanitization...")
test.test_nested_structure_sanitization()
print("✓ Passed")
print("Running test_unicode_non_characters_removed...")
test.test_unicode_non_characters_removed()
print("✓ Passed")
print("Running test_mixed_clean_dirty_data...")
test.test_mixed_clean_dirty_data()
print("✓ Passed")
print("Running test_empty_and_none_strings...")
test.test_empty_and_none_strings()
print("✓ Passed")
print("Running test_specific_surrogate_udc9a...")
test.test_specific_surrogate_udc9a()
print("✓ Passed")
print("Running test_migration_with_surrogate_sanitization...")
test.test_migration_with_surrogate_sanitization()
print("✓ Passed")
print("Running test_empty_values_after_sanitization...")
test.test_empty_values_after_sanitization()
print("✓ Passed")
print("\n✅ All tests passed!")