diff --git a/test_chromadb.py b/test_chromadb.py deleted file mode 100755 index dd4bdd1..0000000 --- a/test_chromadb.py +++ /dev/null @@ -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 的所有分块") - - -if __name__ == "__main__": - main() - diff --git a/test_client.py b/test_client.py deleted file mode 100644 index 8e36dd7..0000000 --- a/test_client.py +++ /dev/null @@ -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() \ No newline at end of file diff --git a/test_scp.py b/test_scp.py deleted file mode 100644 index ea46dc0..0000000 --- a/test_scp.py +++ /dev/null @@ -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) diff --git a/test_sync_configurations.py b/test_sync_configurations.py deleted file mode 100644 index ffe5b2e..0000000 --- a/test_sync_configurations.py +++ /dev/null @@ -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()))