Main function for command-line use. Example usage: python db_loader.py file.txt site_name python db_loader.py https://example.com/feed.rss site_name python db_loader.py data.csv site_name python db_loader.py --delete-site site_name python db_loader.p
()
| 1113 | |
| 1114 | |
| 1115 | async def main(): |
| 1116 | """ |
| 1117 | Main function for command-line use. |
| 1118 | |
| 1119 | Example usage: |
| 1120 | python db_loader.py file.txt site_name |
| 1121 | python db_loader.py https://example.com/feed.rss site_name |
| 1122 | python db_loader.py data.csv site_name |
| 1123 | python db_loader.py --delete-site site_name |
| 1124 | python db_loader.py file.txt site_name --database qdrant_local |
| 1125 | python db_loader.py --force-recompute file.txt site_name |
| 1126 | python db_loader.py --url-list urls.txt site_name |
| 1127 | python db_loader.py --url-list https://example.com/feed_list.txt site_name |
| 1128 | """ |
| 1129 | import argparse |
| 1130 | |
| 1131 | parser = argparse.ArgumentParser(description="Load data into vector database") |
| 1132 | parser.add_argument("--delete-site", action="store_true", |
| 1133 | help="Delete existing entries for the site before loading") |
| 1134 | parser.add_argument("--only-delete", action="store_true", |
| 1135 | help="Only delete entries for the site, don't load data") |
| 1136 | parser.add_argument("--force-recompute", action="store_true", |
| 1137 | help="Force recomputation of embeddings even if a file with embeddings exists") |
| 1138 | parser.add_argument("--url-list", action="store_true", |
| 1139 | help="Treat the input file as a list of URLs to process (one URL per line). The list file itself can be local or a URL.") |
| 1140 | parser.add_argument("--directory", action="store_true", |
| 1141 | help="Treat the input file path as a directory containing files to process.") |
| 1142 | parser.add_argument("file_path", nargs="?", help="Path to the input file or URL or directory containing files to process") |
| 1143 | parser.add_argument("site", help="Site identifier") |
| 1144 | parser.add_argument("--batch-size", type=int, default=100, |
| 1145 | help="Batch size for processing and uploading") |
| 1146 | parser.add_argument("--database", type=str, default=None, |
| 1147 | help="Specific database endpoint to use (from config_retrieval.yaml)") |
| 1148 | |
| 1149 | args = parser.parse_args() |
| 1150 | |
| 1151 | # Validate database if specified |
| 1152 | if args.database and args.database not in CONFIG.retrieval_endpoints: |
| 1153 | parser.error(f"Database endpoint '{args.database}' not found in configuration. Available options: {', '.join(CONFIG.retrieval_endpoints.keys())}") |
| 1154 | |
| 1155 | # Handle delete-only mode |
| 1156 | if args.only_delete: |
| 1157 | await delete_site(args.site, args.database) |
| 1158 | return |
| 1159 | |
| 1160 | # Validate file path if we're not just deleting |
| 1161 | if args.file_path is None and not args.only_delete: |
| 1162 | parser.error("file_path is required unless --only-delete is specified") |
| 1163 | |
| 1164 | # Handle URL list mode |
| 1165 | if args.url_list: |
| 1166 | is_url_path = await is_url(args.file_path) |
| 1167 | if is_url_path: |
| 1168 | print(f"Processing remote URL list from: {args.file_path}") |
| 1169 | else: |
| 1170 | print(f"Processing local URL list file: {args.file_path}") |
| 1171 | |
| 1172 | await loadUrlListToDB(args.file_path, args.site, args.batch_size, args.delete_site, args.force_recompute, args.database) |
no test coverage detected