Spaces:
Sleeping
Sleeping
File size: 1,917 Bytes
da21448 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | 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()
|