删除了测试代码
This commit is contained in:
parent
a6413a76bc
commit
ea1249b977
414
test_chromadb.py
414
test_chromadb.py
|
|
@ -1,414 +0,0 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
ChromaDB 数据查询和测试脚本
|
||||
|
||||
功能：
|
||||
1. 查询 ChromaDB 集合中的文档数量
|
||||
2. 查看数据的存储格式（ids, documents, metadatas, embeddings）
|
||||
3. 展示示例数据
|
||||
4. 统计信息（doc_id 分布、元数据键等）
|
||||
"""
|
||||
|
||||
import sys
|
||||
import json
|
||||
from typing import Dict, List, Any, Optional
|
||||
from collections import Counter
|
||||
import chromadb
|
||||
from chromadb.config import Settings as ChromaSettings
|
||||
from config import settings
|
||||
|
||||
|
||||
def connect_to_chromadb():
|
||||
"""连接到 ChromaDB"""
|
||||
try:
|
||||
if settings.CHROMA_SERVER_HOST:
|
||||
# 使用 HttpClient 模式（Docker/生产环境）
|
||||
print(f"🔗 连接到 ChromaDB 服务器: {settings.CHROMA_SERVER_HOST}:{settings.CHROMA_SERVER_PORT}")
|
||||
client = chromadb.HttpClient(
|
||||
host=settings.CHROMA_SERVER_HOST,
|
||||
port=settings.CHROMA_SERVER_PORT,
|
||||
settings=ChromaSettings(
|
||||
anonymized_telemetry=False,
|
||||
allow_reset=False
|
||||
)
|
||||
)
|
||||
else:
|
||||
# 使用 PersistentClient 模式（本地开发）
|
||||
print(f"🔗 连接到本地 ChromaDB: {settings.CHROMA_DB_PATH}")
|
||||
client = chromadb.PersistentClient(
|
||||
path=settings.CHROMA_DB_PATH,
|
||||
settings=ChromaSettings(
|
||||
anonymized_telemetry=False,
|
||||
allow_reset=False
|
||||
)
|
||||
)
|
||||
|
||||
print("✓ ChromaDB 连接成功\n")
|
||||
return client
|
||||
except Exception as e:
|
||||
print(f"❌ ChromaDB 连接失败: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def get_collection(client, collection_name: str):
|
||||
"""获取或创建集合"""
|
||||
try:
|
||||
collection = client.get_collection(name=collection_name)
|
||||
print(f"✓ 获取集合: {collection_name}\n")
|
||||
return collection
|
||||
except Exception as e:
|
||||
print(f"❌ 获取集合失败: {e}")
|
||||
print(f"提示: 集合 '{collection_name}' 可能不存在")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def get_collection_info(collection) -> Dict[str, Any]:
|
||||
"""获取集合的基本信息"""
|
||||
try:
|
||||
count = collection.count()
|
||||
print(f"📊 集合统计信息:")
|
||||
print(f" - 文档总数: {count}")
|
||||
|
||||
if count == 0:
|
||||
print("\n⚠️ 集合为空，没有数据")
|
||||
return {
|
||||
"count": 0,
|
||||
"sample_ids": [],
|
||||
"sample_documents": [],
|
||||
"sample_metadatas": [],
|
||||
"metadata_keys": [],
|
||||
"doc_id_distribution": {}
|
||||
}
|
||||
|
||||
# 获取所有数据（限制前1000条用于分析）
|
||||
limit = min(1000, count)
|
||||
results = collection.get(limit=limit)
|
||||
|
||||
ids = results.get('ids', [])
|
||||
documents = results.get('documents', [])
|
||||
metadatas = results.get('metadatas', [])
|
||||
|
||||
# 统计 doc_id 分布
|
||||
doc_id_counter = Counter()
|
||||
metadata_keys = set()
|
||||
|
||||
for metadata in metadatas:
|
||||
if metadata:
|
||||
metadata_keys.update(metadata.keys())
|
||||
if 'doc_id' in metadata:
|
||||
doc_id_counter[metadata['doc_id']] += 1
|
||||
|
||||
# 获取示例数据（前5条）
|
||||
sample_size = min(5, len(ids))
|
||||
sample_ids = ids[:sample_size]
|
||||
sample_documents = documents[:sample_size] if documents else []
|
||||
sample_metadatas = metadatas[:sample_size] if metadatas else []
|
||||
|
||||
info = {
|
||||
"count": count,
|
||||
"sample_ids": sample_ids,
|
||||
"sample_documents": sample_documents,
|
||||
"sample_metadatas": sample_metadatas,
|
||||
"metadata_keys": sorted(list(metadata_keys)),
|
||||
"doc_id_distribution": dict(doc_id_counter.most_common(10)) # 前10个最常见的 doc_id
|
||||
}
|
||||
|
||||
return info
|
||||
except Exception as e:
|
||||
print(f"❌ 获取集合信息失败: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def print_data_format(collection_info: Dict[str, Any]):
|
||||
"""打印数据格式说明"""
|
||||
print("\n" + "="*80)
|
||||
print("📋 ChromaDB 数据存储格式")
|
||||
print("="*80)
|
||||
|
||||
if collection_info["count"] == 0:
|
||||
print("\n集合为空，无法展示数据格式")
|
||||
return
|
||||
|
||||
print("\n1. 数据结构说明:")
|
||||
print(" ChromaDB 存储的数据包含以下字段：")
|
||||
print(" - ids: 文档的唯一标识符（字符串列表）")
|
||||
print(" - documents: 文档的文本内容（字符串列表）")
|
||||
print(" - metadatas: 文档的元数据（字典列表）")
|
||||
print(" - embeddings: 文档的向量嵌入（自动生成，不直接显示）")
|
||||
|
||||
print("\n2. ID 格式:")
|
||||
if collection_info["sample_ids"]:
|
||||
print(f" 示例 IDs (前 {len(collection_info['sample_ids'])} 个):")
|
||||
for i, doc_id in enumerate(collection_info["sample_ids"][:5], 1):
|
||||
print(f" [{i}] {doc_id}")
|
||||
# 分析 ID 格式
|
||||
if '_chunk_' in doc_id:
|
||||
parts = doc_id.rsplit('_chunk_', 1)
|
||||
print(f" └─ 格式: {parts[0]}_chunk_{parts[1]} (分块文档)")
|
||||
else:
|
||||
print(f" └─ 格式: {doc_id} (完整文档)")
|
||||
|
||||
print("\n3. 文档内容格式:")
|
||||
print(" 注意: 文档的实际内容存储在 'documents' 字段中，而不是 'metadatas' 中")
|
||||
print(" 'metadatas' 只包含文档的元信息（如 title、doc_id、created_at 等）")
|
||||
if collection_info["sample_documents"]:
|
||||
print(f" 示例文档内容 (前 {len(collection_info['sample_documents'])} 个):")
|
||||
for i, doc_text in enumerate(collection_info["sample_documents"][:3], 1):
|
||||
preview = doc_text[:300] + "..." if len(doc_text) > 300 else doc_text
|
||||
print(f" [{i}] 长度: {len(doc_text)} 字符")
|
||||
# 按行显示，每行缩进
|
||||
preview_lines = preview.split('\n')
|
||||
for line in preview_lines[:5]: # 最多显示5行
|
||||
print(f" {line}")
|
||||
if len(preview_lines) > 5:
|
||||
print(f" ... (还有 {len(preview_lines) - 5} 行)")
|
||||
if len(doc_text) > 300:
|
||||
print(f" ... (还有 {len(doc_text) - 300} 个字符)")
|
||||
|
||||
print("\n4. 元数据格式:")
|
||||
if collection_info["sample_metadatas"]:
|
||||
print(f" 示例元数据 (前 {len(collection_info['sample_metadatas'])} 个):")
|
||||
for i, metadata in enumerate(collection_info["sample_metadatas"][:3], 1):
|
||||
print(f" [{i}] {json.dumps(metadata, ensure_ascii=False, indent=6)}")
|
||||
|
||||
print("\n5. 元数据键列表:")
|
||||
if collection_info["metadata_keys"]:
|
||||
print(f" 所有元数据键 ({len(collection_info['metadata_keys'])} 个):")
|
||||
for key in collection_info["metadata_keys"]:
|
||||
print(f" - {key}")
|
||||
else:
|
||||
print(" 无元数据")
|
||||
|
||||
print("\n6. doc_id 分布 (前10个最常见的文档ID):")
|
||||
if collection_info["doc_id_distribution"]:
|
||||
for doc_id, count in list(collection_info["doc_id_distribution"].items())[:10]:
|
||||
print(f" - {doc_id}: {count} 个分块")
|
||||
else:
|
||||
print(" 无 doc_id 信息")
|
||||
|
||||
|
||||
def print_detailed_sample(collection, limit: int = 3):
|
||||
"""打印详细的示例数据"""
|
||||
print("\n" + "="*80)
|
||||
print(f"📄 详细示例数据 (前 {limit} 条)")
|
||||
print("="*80)
|
||||
|
||||
try:
|
||||
results = collection.get(limit=limit)
|
||||
|
||||
ids = results.get('ids', [])
|
||||
documents = results.get('documents', [])
|
||||
metadatas = results.get('metadatas', [])
|
||||
|
||||
for i in range(len(ids)):
|
||||
print(f"\n{'='*80}")
|
||||
print(f"--- 文档 {i+1} ---")
|
||||
print(f"{'='*80}")
|
||||
print(f"ID: {ids[i]}")
|
||||
|
||||
# 显示文档内容（这是主要的内容）
|
||||
if i < len(documents) and documents[i]:
|
||||
print(f"\n📝 文档内容:")
|
||||
print(f" 长度: {len(documents[i])} 字符")
|
||||
# 显示更多内容（前500字符）
|
||||
doc_preview = documents[i][:500] + "..." if len(documents[i]) > 500 else documents[i]
|
||||
# 按行显示，每行缩进
|
||||
preview_lines = doc_preview.split('\n')
|
||||
for line in preview_lines[:10]: # 最多显示10行
|
||||
print(f" {line}")
|
||||
if len(preview_lines) > 10:
|
||||
print(f" ... (还有 {len(preview_lines) - 10} 行)")
|
||||
if len(documents[i]) > 500:
|
||||
print(f" ... (还有 {len(documents[i]) - 500} 个字符)")
|
||||
else:
|
||||
print(f"\n⚠️ 文档内容为空")
|
||||
|
||||
# 显示元数据
|
||||
if i < len(metadatas) and metadatas[i]:
|
||||
print(f"\n📋 元数据:")
|
||||
for key, value in metadatas[i].items():
|
||||
# 如果值太长，截断显示
|
||||
if isinstance(value, str) and len(value) > 200:
|
||||
value_preview = value[:200] + "..."
|
||||
print(f" - {key}: {value_preview} (长度: {len(value)} 字符)")
|
||||
else:
|
||||
print(f" - {key}: {value}")
|
||||
else:
|
||||
print(f"\n📋 元数据: 无")
|
||||
|
||||
# 尝试获取 embedding 维度（如果可能）
|
||||
try:
|
||||
embedding_info = collection.get(ids=[ids[i]], include=['embeddings'])
|
||||
if embedding_info.get('embeddings') and embedding_info['embeddings']:
|
||||
emb_dim = len(embedding_info['embeddings'][0])
|
||||
print(f"\n🔢 Embedding 维度: {emb_dim}")
|
||||
except Exception:
|
||||
pass # 忽略 embedding 获取错误
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ 获取详细示例数据失败: {e}")
|
||||
|
||||
|
||||
def query_by_doc_id(collection, doc_id: str, limit: int = 1000):
|
||||
"""根据 doc_id 查询所有相关的分块"""
|
||||
print("\n" + "="*80)
|
||||
print(f"🔍 查询 doc_id: {doc_id}")
|
||||
print("="*80)
|
||||
|
||||
try:
|
||||
# 不限制数量，获取所有相关分块
|
||||
results = collection.get(
|
||||
where={"doc_id": doc_id},
|
||||
limit=limit
|
||||
)
|
||||
|
||||
ids = results.get('ids', [])
|
||||
documents = results.get('documents', [])
|
||||
metadatas = results.get('metadatas', [])
|
||||
|
||||
print(f"\n找到 {len(ids)} 个分块:")
|
||||
|
||||
if len(ids) == 0:
|
||||
print(f"⚠️ 未找到 doc_id 为 '{doc_id}' 的文档")
|
||||
return
|
||||
|
||||
# 合并所有分块的内容
|
||||
full_text = ""
|
||||
for i, chunk_id in enumerate(ids):
|
||||
print(f"\n{'='*80}")
|
||||
print(f"--- 分块 {i+1}/{len(ids)} ---")
|
||||
print(f"{'='*80}")
|
||||
print(f"ID: {chunk_id}")
|
||||
|
||||
# 显示文档内容
|
||||
if i < len(documents) and documents[i]:
|
||||
print(f"\n📝 分块内容:")
|
||||
print(f" 长度: {len(documents[i])} 字符")
|
||||
# 显示完整内容或前500字符
|
||||
if len(documents[i]) > 500:
|
||||
doc_preview = documents[i][:500] + "..."
|
||||
print(f" 内容预览 (前500字符):")
|
||||
preview_lines = doc_preview.split('\n')
|
||||
for line in preview_lines[:10]:
|
||||
print(f" {line}")
|
||||
if len(preview_lines) > 10:
|
||||
print(f" ... (还有 {len(preview_lines) - 10} 行)")
|
||||
print(f" ... (还有 {len(documents[i]) - 500} 个字符)")
|
||||
else:
|
||||
preview_lines = documents[i].split('\n')
|
||||
for line in preview_lines[:20]: # 最多显示20行
|
||||
print(f" {line}")
|
||||
if len(preview_lines) > 20:
|
||||
print(f" ... (还有 {len(preview_lines) - 20} 行)")
|
||||
|
||||
full_text += documents[i] + "\n"
|
||||
|
||||
# 显示元数据
|
||||
if i < len(metadatas) and metadatas[i]:
|
||||
print(f"\n📋 元数据:")
|
||||
for key, value in metadatas[i].items():
|
||||
if isinstance(value, str) and len(value) > 200:
|
||||
value_preview = value[:200] + "..."
|
||||
print(f" - {key}: {value_preview} (长度: {len(value)} 字符)")
|
||||
else:
|
||||
print(f" - {key}: {value}")
|
||||
|
||||
# 显示合并后的完整文档
|
||||
if full_text:
|
||||
print(f"\n{'='*80}")
|
||||
print(f"📄 完整文档内容 (所有分块合并)")
|
||||
print(f"{'='*80}")
|
||||
print(f"总长度: {len(full_text)} 字符")
|
||||
print(f"总行数: {len(full_text.split(chr(10)))} 行")
|
||||
print(f"\n内容预览 (前1000字符):")
|
||||
preview = full_text[:1000] + "..." if len(full_text) > 1000 else full_text
|
||||
preview_lines = preview.split('\n')
|
||||
for line in preview_lines[:30]: # 最多显示30行
|
||||
print(f" {line}")
|
||||
if len(preview_lines) > 30:
|
||||
print(f" ... (还有 {len(preview_lines) - 30} 行)")
|
||||
if len(full_text) > 1000:
|
||||
print(f" ... (还有 {len(full_text) - 1000} 个字符)")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ 查询失败: {e}")
|
||||
|
||||
|
||||
def print_statistics(collection_info: Dict[str, Any]):
|
||||
"""打印统计信息"""
|
||||
print("\n" + "="*80)
|
||||
print("📈 统计信息")
|
||||
print("="*80)
|
||||
|
||||
count = collection_info["count"]
|
||||
print(f"\n总文档数: {count}")
|
||||
|
||||
if count == 0:
|
||||
return
|
||||
|
||||
doc_id_dist = collection_info["doc_id_distribution"]
|
||||
if doc_id_dist:
|
||||
unique_docs = len(doc_id_dist)
|
||||
total_chunks = sum(doc_id_dist.values())
|
||||
avg_chunks_per_doc = total_chunks / unique_docs if unique_docs > 0 else 0
|
||||
|
||||
print(f"\n文档统计:")
|
||||
print(f" - 唯一文档数 (doc_id): {unique_docs}")
|
||||
print(f" - 总分块数: {total_chunks}")
|
||||
print(f" - 平均每个文档的分块数: {avg_chunks_per_doc:.2f}")
|
||||
|
||||
# 分块数分布
|
||||
chunk_counts = Counter(doc_id_dist.values())
|
||||
print(f"\n分块数分布:")
|
||||
for chunk_count, doc_count in sorted(chunk_counts.items()):
|
||||
print(f" - {chunk_count} 个分块: {doc_count} 个文档")
|
||||
|
||||
metadata_keys = collection_info["metadata_keys"]
|
||||
print(f"\n元数据键数量: {len(metadata_keys)}")
|
||||
if metadata_keys:
|
||||
print(f"元数据键列表: {', '.join(metadata_keys)}")
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
print("="*80)
|
||||
print("ChromaDB 数据查询工具")
|
||||
print("="*80)
|
||||
print()
|
||||
|
||||
# 连接 ChromaDB
|
||||
client = connect_to_chromadb()
|
||||
|
||||
# 获取集合
|
||||
collection_name = settings.CHROMA_COLLECTION_NAME
|
||||
collection = get_collection(client, collection_name)
|
||||
|
||||
# 获取集合信息
|
||||
collection_info = get_collection_info(collection)
|
||||
|
||||
# 打印数据格式
|
||||
print_data_format(collection_info)
|
||||
|
||||
# 打印详细示例
|
||||
if collection_info["count"] > 0:
|
||||
print_detailed_sample(collection, limit=3)
|
||||
|
||||
# 打印统计信息
|
||||
print_statistics(collection_info)
|
||||
|
||||
# 如果提供了 doc_id 参数，查询该文档
|
||||
if len(sys.argv) > 1:
|
||||
doc_id = sys.argv[1]
|
||||
query_by_doc_id(collection, doc_id)
|
||||
|
||||
print("\n" + "="*80)
|
||||
print("✓ 查询完成")
|
||||
print("="*80)
|
||||
print("\n使用说明:")
|
||||
print(" python test_chromadb.py # 查看所有数据")
|
||||
print(" python test_chromadb.py <doc_id> # 查询指定 doc_id 的所有分块")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
226
test_client.py
226
test_client.py
|
|
@ -1,226 +0,0 @@
|
|||
"""
|
||||
Test client for RAG API
|
||||
"""
|
||||
import requests
|
||||
import json
|
||||
import sys
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
|
||||
class ServiceChecker:
|
||||
"""Service dependency checker"""
|
||||
|
||||
@staticmethod
|
||||
def check_api_service(exit_on_failure=False):
|
||||
"""Check if API service is running
|
||||
|
||||
Args:
|
||||
exit_on_failure: If True, exit the program if API service is not running
|
||||
"""
|
||||
print("Checking API service...")
|
||||
|
||||
# First check if API service is already running
|
||||
try:
|
||||
response = requests.get("http://localhost:8001/health", timeout=10)
|
||||
if response.status_code == 200:
|
||||
print("✓ API service is running")
|
||||
return True
|
||||
else:
|
||||
print(f"✗ API service returned status code: {response.status_code}")
|
||||
if exit_on_failure:
|
||||
print("API service is not running. Exiting program.")
|
||||
sys.exit(1)
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"✗ API service is not running: {e}")
|
||||
|
||||
# Try to start API service using Docker Compose
|
||||
print("Attempting to start API service using Docker Compose...")
|
||||
try:
|
||||
import os
|
||||
|
||||
# Ask user if they want to rebuild the image
|
||||
rebuild = input("Do you want to rebuild the Docker image before starting services? (y/n, default: n): ").strip().lower() == 'y'
|
||||
|
||||
if rebuild:
|
||||
# Check if rebuild script exists
|
||||
if os.path.exists('./rebuild_and_start.sh'):
|
||||
print("Rebuilding and starting services using rebuild_and_start.sh...")
|
||||
result = subprocess.run(['bash', './rebuild_and_start.sh'], check=False, text=True)
|
||||
if result.returncode != 0:
|
||||
print(f"✗ rebuild_and_start.sh failed with exit code {result.returncode}")
|
||||
raise subprocess.CalledProcessError(result.returncode, result.args)
|
||||
else:
|
||||
print("Warning: rebuild_and_start.sh not found. Using regular docker-compose build and up instead.")
|
||||
|
||||
# Build the image with detailed error reporting
|
||||
print("Building Docker image...")
|
||||
result = subprocess.run(['docker-compose', 'build'], check=False, capture_output=True, text=True)
|
||||
if result.returncode != 0:
|
||||
print(f"✗ Docker Compose build failed with exit code {result.returncode}")
|
||||
print(f"Stdout: {result.stdout}")
|
||||
print(f"Stderr: {result.stderr}")
|
||||
raise subprocess.CalledProcessError(result.returncode, result.args, result.stdout, result.stderr)
|
||||
print("✓ Docker image built successfully")
|
||||
|
||||
# Start services with detailed error reporting
|
||||
print("Starting Docker Compose services...")
|
||||
result = subprocess.run(['docker-compose', 'up', '-d'], check=False, capture_output=True, text=True)
|
||||
if result.returncode != 0:
|
||||
print(f"✗ Docker Compose up failed with exit code {result.returncode}")
|
||||
print(f"Stdout: {result.stdout}")
|
||||
print(f"Stderr: {result.stderr}")
|
||||
raise subprocess.CalledProcessError(result.returncode, result.args, result.stdout, result.stderr)
|
||||
print("✓ Docker Compose services started successfully")
|
||||
else:
|
||||
print("Starting services without rebuilding...")
|
||||
result = subprocess.run(
|
||||
["docker-compose", "up", "-d"],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True
|
||||
)
|
||||
if result.returncode == 0:
|
||||
print("✓ Docker Compose services started successfully")
|
||||
else:
|
||||
print(f"✗ Docker Compose up failed with exit code {result.returncode}")
|
||||
print(f"Stdout: {result.stdout}")
|
||||
print(f"Stderr: {result.stderr}")
|
||||
raise subprocess.CalledProcessError(result.returncode, result.args, result.stdout, result.stderr)
|
||||
|
||||
# Wait for the API service to initialize with a retry loop
|
||||
print("Waiting for API service to initialize...")
|
||||
max_retries = 10
|
||||
retry_interval = 5 # seconds between retries
|
||||
|
||||
for attempt in range(1, max_retries + 1):
|
||||
try:
|
||||
print(f" Health check attempt {attempt}/{max_retries}...")
|
||||
response = requests.get("http://localhost:8001/health", timeout=10)
|
||||
if response.status_code == 200:
|
||||
print("✓ API service is now running")
|
||||
return True
|
||||
except requests.exceptions.ConnectionError:
|
||||
pass # Service not ready yet, continue retrying
|
||||
except Exception as e:
|
||||
print(f" Health check failed with unexpected error: {e}")
|
||||
|
||||
# Wait before next attempt
|
||||
if attempt < max_retries:
|
||||
time.sleep(retry_interval)
|
||||
|
||||
# All attempts failed
|
||||
print(f"✗ API service still not running after {max_retries} attempts")
|
||||
if exit_on_failure:
|
||||
print("Failed to start API service. Exiting program.")
|
||||
sys.exit(1)
|
||||
return False
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"✗ Failed to start Docker Compose services: {e}")
|
||||
if exit_on_failure:
|
||||
print("Failed to start API service. Exiting program.")
|
||||
sys.exit(1)
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"✗ Unexpected error when starting services: {e}")
|
||||
if exit_on_failure:
|
||||
print("Failed to start API service. Exiting program.")
|
||||
sys.exit(1)
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def check_all_services(exit_on_failure=False):
|
||||
"""Check all required services and start them if needed
|
||||
|
||||
Args:
|
||||
exit_on_failure: If True, exit the program if any critical service fails
|
||||
"""
|
||||
print("Checking all required services...")
|
||||
|
||||
# Check API service first (critical)
|
||||
if not ServiceChecker.check_api_service(exit_on_failure=exit_on_failure):
|
||||
return False
|
||||
|
||||
print("\n✅ All critical services are running!")
|
||||
return True
|
||||
|
||||
|
||||
def test_health():
|
||||
"""Test health check endpoint"""
|
||||
print("Testing health check...")
|
||||
try:
|
||||
response = requests.get("http://localhost:8001/health", timeout=10)
|
||||
print(f"Status: {response.status_code}")
|
||||
print(f"Response: {response.json()}")
|
||||
return response.status_code == 200
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def test_query(query: str):
|
||||
"""Test query endpoint"""
|
||||
print(f"Testing query: {query}")
|
||||
try:
|
||||
url = "http://localhost:8001/query"
|
||||
data = {
|
||||
"query": query,
|
||||
"stream": True,
|
||||
"top_k": 5
|
||||
}
|
||||
response = requests.post(url, json=data, stream=True, timeout=(10, 300))
|
||||
if response.status_code != 200:
|
||||
print(f"Error: {response.status_code}")
|
||||
try:
|
||||
print(f"Error message: {response.text}")
|
||||
except:
|
||||
pass
|
||||
return False
|
||||
|
||||
print("Response (streaming):")
|
||||
# Process streaming response
|
||||
try:
|
||||
for chunk in response.iter_content(chunk_size=1024, decode_unicode=True):
|
||||
if chunk:
|
||||
print(chunk, end='', flush=True)
|
||||
print() # Add a newline at the end
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"\nStreaming error: {e}")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
"""Run all tests"""
|
||||
print("=" * 50)
|
||||
print("RAG API Test Client")
|
||||
print("=" * 50)
|
||||
|
||||
# Check all services with exit_on_failure=True
|
||||
if not ServiceChecker.check_all_services(exit_on_failure=True):
|
||||
sys.exit(1)
|
||||
|
||||
# Test health endpoint
|
||||
if not test_health():
|
||||
print("\n❌ Health check failed")
|
||||
sys.exit(1)
|
||||
|
||||
# Test queries
|
||||
test_queries = [
|
||||
"这个系统是如何工作的？"
|
||||
]
|
||||
|
||||
for query in test_queries:
|
||||
test_query(query)
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("Tests completed!")
|
||||
print("=" * 50)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
144
test_scp.py
144
test_scp.py
|
|
@ -1,144 +0,0 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script to verify SCP folder access functionality
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
from pathlib import Path
|
||||
import paramiko
|
||||
from loguru import logger
|
||||
|
||||
# Add current directory to Python path
|
||||
sys.path.append('.')
|
||||
|
||||
from config import settings
|
||||
|
||||
def test_scp_folder_access(host, port, username, password=None, key_path=None, remote_folder_path):
|
||||
"""
|
||||
Test SCP folder access functionality
|
||||
|
||||
Args:
|
||||
host: Remote host address
|
||||
port: SSH port
|
||||
username: SSH username
|
||||
password: SSH password (optional)
|
||||
key_path: SSH private key path (optional)
|
||||
remote_folder_path: Remote folder path
|
||||
|
||||
Returns:
|
||||
dict: Test results
|
||||
"""
|
||||
results = {
|
||||
"success": False,
|
||||
"message": "",
|
||||
"files": [],
|
||||
"folders": []
|
||||
}
|
||||
|
||||
try:
|
||||
logger.info(f"Testing SCP access to {host}:{port}...")
|
||||
|
||||
# Create SSH client
|
||||
ssh_client = paramiko.SSHClient()
|
||||
ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||||
|
||||
# Connect to remote host
|
||||
connect_kwargs = {
|
||||
'hostname': host,
|
||||
'port': port,
|
||||
'username': username
|
||||
}
|
||||
|
||||
if key_path:
|
||||
connect_kwargs['key_filename'] = key_path
|
||||
elif password:
|
||||
connect_kwargs['password'] = password
|
||||
else:
|
||||
raise ValueError("Either password or key_path must be provided for SCP")
|
||||
|
||||
ssh_client.connect(**connect_kwargs)
|
||||
logger.info(f"✓ Connected to {host}")
|
||||
|
||||
# Create SCP client
|
||||
scp_client = ssh_client.open_sftp()
|
||||
logger.info("✓ SFTP client created")
|
||||
|
||||
# List files in remote folder
|
||||
logger.info(f"Listing files in {remote_folder_path}...")
|
||||
try:
|
||||
items = scp_client.listdir_attr(remote_folder_path)
|
||||
|
||||
for item in items:
|
||||
if item.filename in ('.', '..'):
|
||||
continue
|
||||
|
||||
item_path = os.path.join(remote_folder_path, item.filename)
|
||||
if item.st_mode & 0o170000 == 0o040000: # Directory
|
||||
results["folders"].append(item_path)
|
||||
else: # File
|
||||
results["files"].append(item_path)
|
||||
|
||||
logger.info(f"✓ Found {len(results['files'])} files and {len(results['folders'])} folders")
|
||||
results["success"] = True
|
||||
results["message"] = f"Successfully accessed remote folder with {len(results['files'])} files and {len(results['folders'])} folders"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing remote folder: {e}")
|
||||
results["message"] = f"Error listing remote folder: {str(e)}"
|
||||
|
||||
# Close connections
|
||||
scp_client.close()
|
||||
ssh_client.close()
|
||||
logger.info("✓ Connections closed")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"SCP test failed: {e}")
|
||||
results["message"] = f"SCP test failed: {str(e)}"
|
||||
|
||||
return results
|
||||
|
||||
if __name__ == "__main__":
|
||||
logger.info("=== SCP Folder Access Test ===")
|
||||
|
||||
# Get test parameters from command line
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Test SCP folder access functionality")
|
||||
parser.add_argument("--host", required=True, help="Remote host address")
|
||||
parser.add_argument("--port", type=int, default=22, help="SSH port")
|
||||
parser.add_argument("--username", required=True, help="SSH username")
|
||||
parser.add_argument("--password", help="SSH password (optional)")
|
||||
parser.add_argument("--key-path", help="SSH private key path (optional)")
|
||||
parser.add_argument("--folder", required=True, help="Remote folder path")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Run test
|
||||
results = test_scp_folder_access(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
username=args.username,
|
||||
password=args.password,
|
||||
key_path=args.key_path,
|
||||
remote_folder_path=args.folder
|
||||
)
|
||||
|
||||
logger.info("\n=== Test Results ===")
|
||||
logger.info(f"Success: {results['success']}")
|
||||
logger.info(f"Message: {results['message']}")
|
||||
|
||||
if results['success']:
|
||||
logger.info(f"Files found: {len(results['files'])}")
|
||||
for file in results['files'][:10]: # Show first 10 files
|
||||
logger.info(f" - {file}")
|
||||
if len(results['files']) > 10:
|
||||
logger.info(f" ... and {len(results['files']) - 10} more files")
|
||||
|
||||
logger.info(f"Folders found: {len(results['folders'])}")
|
||||
for folder in results['folders'][:10]: # Show first 10 folders
|
||||
logger.info(f" - {folder}")
|
||||
if len(results['folders']) > 10:
|
||||
logger.info(f" ... and {len(results['folders']) - 10} more folders")
|
||||
|
||||
sys.exit(0 if results['success'] else 1)
|
||||
|
|
@ -1,338 +0,0 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script for testing document sync with various database configurations
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
import argparse
|
||||
|
||||
# Add current directory to Python path
|
||||
sys.path.append('.')
|
||||
|
||||
from config import Settings, DatabaseDataSourceConfig
|
||||
from sync_service import SyncService
|
||||
|
||||
|
||||
class ConfigurationLoader:
|
||||
"""Load configurations from database_config.json"""
|
||||
|
||||
def __init__(self, config_file_path: str = "./database_config.json"):
|
||||
"""
|
||||
Initialize configuration loader
|
||||
|
||||
Args:
|
||||
config_file_path: Path to the database configuration file
|
||||
"""
|
||||
self.config_file_path = config_file_path
|
||||
self.configurations = []
|
||||
|
||||
# Load configurations from file
|
||||
self._load_configurations()
|
||||
|
||||
def _load_configurations(self):
|
||||
"""Load configurations from the database_config.json file"""
|
||||
try:
|
||||
with open(self.config_file_path, "r", encoding="utf-8") as f:
|
||||
self.configurations = json.load(f)
|
||||
|
||||
print(f"Loaded {len(self.configurations)} configurations from {self.config_file_path}")
|
||||
|
||||
# Print summary of available configurations
|
||||
print("\nAvailable configurations:")
|
||||
for i, config in enumerate(self.configurations, 1):
|
||||
print(f"{i}. {config['name']} (database: {config['database']}, table: {config['table_name']})")
|
||||
print(f" Content columns: {config['content_column']}")
|
||||
if config.get('file_source_type'):
|
||||
print(f" File source: {config['file_source_type']}")
|
||||
print()
|
||||
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to load configurations from {self.config_file_path}: {e}")
|
||||
|
||||
def get_configuration_by_name(self, config_name: str):
|
||||
"""
|
||||
Get configuration by name
|
||||
|
||||
Args:
|
||||
config_name: Name of the configuration to retrieve
|
||||
|
||||
Returns:
|
||||
Configuration dictionary
|
||||
"""
|
||||
for config in self.configurations:
|
||||
if config['name'] == config_name:
|
||||
return config
|
||||
raise ValueError(f"Configuration with name '{config_name}' not found")
|
||||
|
||||
def get_configurations_by_type(self, config_type: str = "database"):
|
||||
"""
|
||||
Get configurations by type
|
||||
|
||||
Args:
|
||||
config_type: Type of configurations to retrieve
|
||||
|
||||
Returns:
|
||||
List of configuration dictionaries
|
||||
"""
|
||||
return [config for config in self.configurations if config['type'] == config_type]
|
||||
|
||||
def get_all_configurations(self):
|
||||
"""
|
||||
Get all configurations
|
||||
|
||||
Returns:
|
||||
List of all configuration dictionaries
|
||||
"""
|
||||
return self.configurations
|
||||
|
||||
def create_test_config_file(self, configurations):
|
||||
"""
|
||||
Create a temporary test configuration file
|
||||
|
||||
Args:
|
||||
configurations: List of configurations to include in the test file
|
||||
|
||||
Returns:
|
||||
Path to the created test configuration file
|
||||
"""
|
||||
config_path = f"./test_config_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
|
||||
|
||||
with open(config_path, "w", encoding="utf-8") as f:
|
||||
json.dump(configurations, f, ensure_ascii=False, indent=2)
|
||||
|
||||
print(f"Created test configuration file: {config_path}")
|
||||
return config_path
|
||||
|
||||
|
||||
async def test_sync(configurations, sync_type: str = "all", force: bool = False):
|
||||
"""
|
||||
Test document sync with the specified configurations
|
||||
|
||||
Args:
|
||||
configurations: List of configurations to test
|
||||
sync_type: Type of sync to perform ("all" for full sync, "incremental" for incremental sync)
|
||||
force: Whether to force full sync (ignored for incremental)
|
||||
"""
|
||||
if not configurations:
|
||||
print("No configurations to test")
|
||||
return False
|
||||
|
||||
config_names = [config["name"] for config in configurations]
|
||||
print(f"\n=== Testing {sync_type.upper()} Sync with Configuration(s): {', '.join(config_names)} ===")
|
||||
|
||||
loader = ConfigurationLoader()
|
||||
|
||||
# Define database path outside try block for cleanup
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
import json
|
||||
|
||||
DATA_DIR = Path(__file__).parent / "data"
|
||||
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||||
DB_PATH = DATA_DIR / "sessions.db"
|
||||
|
||||
try:
|
||||
# Write test configurations to SQLite database
|
||||
conn = sqlite3.connect(DB_PATH)
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Create table if not exists
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS data_sources (
|
||||
name TEXT PRIMARY KEY,
|
||||
config TEXT,
|
||||
update_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
''')
|
||||
|
||||
# Insert/update test configurations
|
||||
for config in configurations:
|
||||
config_json = json.dumps(config, ensure_ascii=False, indent=2)
|
||||
cursor.execute('''
|
||||
INSERT OR REPLACE INTO data_sources (name, config, update_at)
|
||||
VALUES (?, ?, CURRENT_TIMESTAMP)
|
||||
''', (config['name'], config_json))
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
print(f"Inserted {len(configurations)} test configurations into SQLite database")
|
||||
|
||||
# Initialize settings normally (will read from SQLite)
|
||||
from config import Settings
|
||||
settings = Settings()
|
||||
|
||||
# Initialize sync service
|
||||
print("\nInitializing SyncService...")
|
||||
sync_service = SyncService()
|
||||
|
||||
print(f"\nConfiguration details:")
|
||||
for db_config in sync_service.db_configs:
|
||||
print(f"- Name: {db_config.name}")
|
||||
print(f"- Database: {db_config.database}")
|
||||
print(f"- Table: {db_config.table_name}")
|
||||
print(f"- Content columns: {db_config.content_columns}")
|
||||
if db_config.file_source_type:
|
||||
print(f"- File source: {db_config.file_source_type}")
|
||||
if db_config.file_system_base_path:
|
||||
print(f"- File base path: {db_config.file_system_base_path}")
|
||||
if db_config.updated_at_column:
|
||||
print(f"- Updated at column: {db_config.updated_at_column}")
|
||||
print()
|
||||
|
||||
# Perform sync
|
||||
start_time = datetime.now()
|
||||
|
||||
if sync_type == "all":
|
||||
print(f"Starting full sync (force={force})...")
|
||||
await sync_service.sync_all(force=force)
|
||||
elif sync_type == "incremental":
|
||||
print("Starting incremental sync...")
|
||||
await sync_service.sync_incremental()
|
||||
else:
|
||||
raise ValueError(f"Unknown sync type: {sync_type}")
|
||||
|
||||
duration = (datetime.now() - start_time).total_seconds()
|
||||
print(f"\n✓ Sync completed successfully in {duration:.2f} seconds")
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n✗ Sync failed: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
finally:
|
||||
# Clean up: remove test configurations from SQLite database
|
||||
try:
|
||||
conn = sqlite3.connect(DB_PATH)
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Delete test configurations
|
||||
for config in configurations:
|
||||
cursor.execute('DELETE FROM data_sources WHERE name = ?', (config['name'],))
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
print(f"Removed {len(configurations)} test configurations from SQLite database")
|
||||
except Exception as e:
|
||||
print(f"Error during cleanup: {e}")
|
||||
|
||||
|
||||
async def main():
|
||||
"""
|
||||
Main function to run tests
|
||||
"""
|
||||
parser = argparse.ArgumentParser(description="Test document sync with configurations from database_config.json")
|
||||
|
||||
# Add mutually exclusive group for configuration selection
|
||||
group = parser.add_mutually_exclusive_group()
|
||||
|
||||
group.add_argument(
|
||||
"--all",
|
||||
action="store_true",
|
||||
help="Test all configurations"
|
||||
)
|
||||
|
||||
group.add_argument(
|
||||
"--name",
|
||||
type=str,
|
||||
help="Test specific configuration by name"
|
||||
)
|
||||
|
||||
group.add_argument(
|
||||
"--list",
|
||||
action="store_true",
|
||||
help="List all available configurations and exit"
|
||||
)
|
||||
|
||||
# Sync options
|
||||
parser.add_argument(
|
||||
"--sync-type",
|
||||
type=str,
|
||||
choices=["all", "incremental"],
|
||||
default="all",
|
||||
help="Type of sync to perform (default: all)"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--force",
|
||||
action="store_true",
|
||||
help="Force full sync (ignored for incremental sync)"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
print("=" * 70)
|
||||
print("RAG Sync Configuration Tester")
|
||||
print("=" * 70)
|
||||
|
||||
# Initialize configuration loader
|
||||
loader = ConfigurationLoader()
|
||||
|
||||
# List configurations and exit if requested
|
||||
if args.list:
|
||||
print("\nConfiguration listing complete.")
|
||||
return 0
|
||||
|
||||
# Determine which configurations to test
|
||||
configurations_to_test = []
|
||||
|
||||
if args.all:
|
||||
# Test all configurations
|
||||
configurations_to_test = loader.get_all_configurations()
|
||||
print(f"\nSelected all {len(configurations_to_test)} configurations for testing")
|
||||
|
||||
elif args.name:
|
||||
# Test specific configuration by name
|
||||
try:
|
||||
config = loader.get_configuration_by_name(args.name)
|
||||
configurations_to_test = [config]
|
||||
print(f"\nSelected configuration: {args.name}")
|
||||
except ValueError as e:
|
||||
print(f"\nError: {e}")
|
||||
return 1
|
||||
|
||||
else:
|
||||
# Default: ask user to select configuration
|
||||
print("\nPlease select configuration(s) to test (comma-separated numbers or 'all'):")
|
||||
print("Example: 1,3 or all")
|
||||
|
||||
user_input = input("Selection: ").strip()
|
||||
|
||||
if user_input.lower() == "all":
|
||||
configurations_to_test = loader.get_all_configurations()
|
||||
else:
|
||||
try:
|
||||
indices = [int(idx.strip()) - 1 for idx in user_input.split(",")]
|
||||
configurations_to_test = [loader.configurations[i] for i in indices]
|
||||
except (ValueError, IndexError) as e:
|
||||
print(f"\nInvalid selection: {e}")
|
||||
return 1
|
||||
|
||||
# Test the selected configurations
|
||||
if configurations_to_test:
|
||||
print(f"\nTesting {len(configurations_to_test)} configuration(s)")
|
||||
|
||||
# Test configurations
|
||||
success = await test_sync(configurations_to_test, args.sync_type, args.force)
|
||||
|
||||
# Print results
|
||||
print("\n" + "=" * 70)
|
||||
if success:
|
||||
print("🎉 All selected configurations passed the sync test!")
|
||||
return 0
|
||||
else:
|
||||
print("❌ Some configurations failed the sync test.")
|
||||
return 1
|
||||
else:
|
||||
print("\nNo configurations selected for testing.")
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(asyncio.run(main()))
|
||||
Loading…
Reference in New Issue