Spaces:
Sleeping
Sleeping
| import argparse | |
| import logging | |
| import time | |
| from pathlib import Path | |
| from neo4j.exceptions import DriverError, Neo4jError | |
| from techdocs_graphrag.ingestion import load_documents | |
| from techdocs_graphrag.services import Services, build_services | |
| RETRY_DELAY_SECONDS = 2 | |
| def parse_arguments() -> argparse.Namespace: | |
| parser = argparse.ArgumentParser(description="Indexe un corpus Markdown dans Neo4j.") | |
| parser.add_argument("data_dir", type=Path, nargs="?", default=Path("data")) | |
| parser.add_argument("--reset", action="store_true", help="Vide la base avant indexation.") | |
| parser.add_argument("--if-empty", action="store_true", help="N'indexe que si la base est vide.") | |
| parser.add_argument("--wait", type=int, default=0, metavar="SECONDS", help="Attend Neo4j.") | |
| return parser.parse_args() | |
| def wait_for_database(services: Services, timeout_seconds: int) -> None: | |
| deadline = time.monotonic() + timeout_seconds | |
| while True: | |
| try: | |
| services.store.verify_connectivity() | |
| return | |
| except (DriverError, Neo4jError): | |
| if time.monotonic() >= deadline: | |
| raise | |
| time.sleep(RETRY_DELAY_SECONDS) | |
| def main() -> None: | |
| logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") | |
| arguments = parse_arguments() | |
| services = build_services() | |
| try: | |
| wait_for_database(services, arguments.wait) | |
| if arguments.if_empty and services.store.stats()["chunks"]: | |
| print("Base déjà indexée.") | |
| return | |
| if arguments.reset: | |
| services.store.reset() | |
| report = services.indexer.run( | |
| load_documents(arguments.data_dir), | |
| on_progress=lambda done, total: print(f"\r{done}/{total} chunks", end="", flush=True), | |
| ) | |
| print(f"\n{report}") | |
| finally: | |
| services.close() | |
| if __name__ == "__main__": | |
| main() | |