Coverage for app \ vector_store \ qdrant_store.py: 100%

1from typing import Optional 

2 

3from qdrant_client import QdrantClient 

4from qdrant_client.models import VectorParams, Distance, PayloadSchemaType 

5from qdrant_client.http.exceptions import UnexpectedResponse 

6 

7from app.config import settings 

8from app.utils.logger import get_logger 

9 

10logger = get_logger(__name__) 

11 

12COLLECTION = settings.QDRANT_COLLECTION 

13 

14 

15def get_client() -> QdrantClient: 

16 """ 

17 Create and return a Qdrant client. 

18 Supports local, Docker, and cloud-based Qdrant. 

19 """ 

20 try: 

21 client = QdrantClient( 

22 url=settings.QDRANT_URL, 

23 api_key=settings.QDRANT_API_KEY, 

24 timeout=settings.QDRANT_TIMEOUT, 

25 ) 

26 logger.info("Qdrant client initialized") 

27 return client 

28 except Exception as e: 

29 logger.exception("Failed to initialize Qdrant client") 

30 raise RuntimeError("Qdrant connection failed") from e 

31 

32 

33def create_collection_if_not_exists(client: QdrantClient) -> None: 

34 """ 

35 Create the vector collection if it does not exist. 

36 Safe to call multiple times. 

37 """ 

38 try: 

39 existing = {c.name for c in client.get_collections().collections} 

40 

41 if COLLECTION in existing: 

42 logger.info("Qdrant collection already exists", extra={"collection": COLLECTION}) 

43 return 

44 

45 client.create_collection( 

46 collection_name=COLLECTION, 

47 vectors_config=VectorParams( 

48 size=settings.EMBEDDING_DIM, 

49 distance=Distance.COSINE, 

50 ), 

51 ) 

52 logger.info("Created Qdrant collection", extra={"collection": COLLECTION}) 

53 

54 except Exception as e: 

55 logger.exception("Failed to create Qdrant collection") 

56 raise 

57 

58 

59def _create_payload_index_safe( 

60 client: QdrantClient, 

61 field_name: str, 

62 field_schema: PayloadSchemaType, 

63) -> None: 

64 """ 

65 Create payload index safely (idempotent). 

66 """ 

67 try: 

68 client.create_payload_index( 

69 collection_name=COLLECTION, 

70 field_name=field_name, 

71 field_schema=field_schema, 

72 ) 

73 logger.info("Created payload index", extra={"field": field_name}) 

74 except UnexpectedResponse as e: 

75 # Index probably already exists 

76 logger.debug("Payload index already exists", extra={"field": field_name}) 

77 except Exception: 

78 logger.exception("Failed to create payload index", extra={"field": field_name}) 

79 raise 

80 

81 

82def create_indexes(client: QdrantClient) -> None: 

83 """ 

84 Create payload indexes for efficient filtering. 

85 Safe to call multiple times. 

86 """ 

87 _create_payload_index_safe(client, "pmid", PayloadSchemaType.KEYWORD) 

88 _create_payload_index_safe(client, "year", PayloadSchemaType.INTEGER) 

89 _create_payload_index_safe(client, "journal", PayloadSchemaType.KEYWORD) 

90 _create_payload_index_safe(client, "study_type", PayloadSchemaType.KEYWORD) 

91 _create_payload_index_safe(client, "entities.drugs", PayloadSchemaType.KEYWORD) 

92 

93 logger.info("Qdrant payload indexes ensured")