Coverage for src/qdrant_loader_core/graph/__init__.py: 92%

1from __future__ import annotations 

2 

3import asyncio 

4from typing import TYPE_CHECKING 

5 

6from .base import GraphEdge, GraphNode, GraphStore, SubGraph 

7from .schema.utils import init_schema 

8 

9if TYPE_CHECKING: 

10 from .falkor_store import FalkorGraphStore 

11 

12try: 

13 from .falkor_store import FalkorGraphStore 

14except ImportError: 

15 FalkorGraphStore = None 

16 

17 

18__all__ = [ 

19 "FalkorGraphStore", 

20 "GraphNode", 

21 "GraphEdge", 

22 "GraphStore", 

23 "SubGraph", 

24] 

25 

26 

27_graph_store: FalkorGraphStore | None = None 

28_graph_store_lock = asyncio.Lock() 

29 

30 

31async def get_graph_store( 

32 host: str | None = None, 

33 port: int | None = None, 

34 password: str | None = None, 

35 graph_name: str | None = None, 

36 max_connections: int | None = None, 

37) -> FalkorGraphStore: 

38 global _graph_store 

39 

40 if FalkorGraphStore is None: 

41 raise ModuleNotFoundError( 

42 "FalkorGraphStore requires the 'graph' extra.\n" 

43 "Install it with:\n" 

44 " pip install qdrant-loader-core[graph]" 

45 ) 

46 

47 if _graph_store is None: 

48 async with _graph_store_lock: 

49 if _graph_store is None: 

50 final_host = host or "localhost" 

51 final_port = port if port is not None else 6379 

52 final_graph = graph_name or "default_graph" 

53 final_max_conn = max_connections if max_connections is not None else 10 

54 

55 _graph_store = FalkorGraphStore( 

56 host=final_host, 

57 port=int(final_port), 

58 password=password, 

59 graph_name=final_graph, 

60 max_connections=final_max_conn, 

61 ) 

62 

63 await init_schema(_graph_store) 

64 

65 return _graph_store