| 160 | |
| 161 | |
| 162 | def main(): |
| 163 | parser = argparse.ArgumentParser( |
| 164 | description="Load base vectors into the vector database service via /bulk_insert." |
| 165 | ) |
| 166 | parser.add_argument( |
| 167 | "--server-url", |
| 168 | default="http://127.0.0.1:8080", |
| 169 | help="Base URL of the vector database server (default: http://127.0.0.1:8080)", |
| 170 | ) |
| 171 | parser.add_argument( |
| 172 | "--data-dir", |
| 173 | default=os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data"), |
| 174 | help="Directory containing base vector JSON files (default: data/)", |
| 175 | ) |
| 176 | parser.add_argument( |
| 177 | "--batch-size", |
| 178 | type=int, |
| 179 | default=5000, |
| 180 | help="Number of vectors per bulk_insert request (default: 5000)", |
| 181 | ) |
| 182 | args = parser.parse_args() |
| 183 | |
| 184 | if args.batch_size <= 0: |
| 185 | print("Error: --batch-size must be a positive integer.", file=sys.stderr) |
| 186 | sys.exit(1) |
| 187 | |
| 188 | try: |
| 189 | load_data(args.server_url, args.data_dir, args.batch_size) |
| 190 | except FileNotFoundError as e: |
| 191 | print(f"Error: {e}", file=sys.stderr) |
| 192 | sys.exit(1) |
| 193 | except RuntimeError as e: |
| 194 | print(f"Error: {e}", file=sys.stderr) |
| 195 | sys.exit(1) |
| 196 | |
| 197 | |
| 198 | if __name__ == "__main__": |