FazBrowse GitHub Viewer | Trending | | Home

Tools: [Download Repo ZIP] [View Raw Code] [Original HTTPS Page]

testcontainers-python/modules/qdrant/example_basic.py at main · f18m/testcontainers-python · GitHub

149 lines (128 loc) · 5.26 KB

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

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

import json

from datetime import datetime

import numpy as np

from qdrant_client import QdrantClient

from qdrant_client.http import models

from testcontainers.qdrant import QdrantContainer

def basic_example():

with QdrantContainer() as qdrant:

# Get connection parameters

host = qdrant.get_container_host_ip()

port = qdrant.get_exposed_port(qdrant.port)

# Create Qdrant client

client = QdrantClient(host=host, port=port)

print("Connected to Qdrant")

# Create collection

collection_name = "test_collection"

vector_size = 128

client.create_collection(

collection_name=collection_name,

vectors_config=models.VectorParams(size=vector_size, distance=models.Distance.COSINE),

)

print(f"Created collection: {collection_name}")

# Generate test vectors and payloads

num_vectors = 5

vectors = np.random.rand(num_vectors, vector_size).tolist()

payloads = [

{

"text": "AI and machine learning are transforming industries",

"category": "Technology",

"tags": ["AI", "ML", "innovation"],

"timestamp": datetime.utcnow().isoformat(),

},

{

"text": "New study reveals benefits of meditation",

"category": "Health",

"tags": ["wellness", "mental health"],

"timestamp": datetime.utcnow().isoformat(),

},

{

"text": "Global warming reaches critical levels",

"category": "Environment",

"tags": ["climate", "sustainability"],

"timestamp": datetime.utcnow().isoformat(),

},

{

"text": "Stock market shows strong growth",

"category": "Finance",

"tags": ["investing", "markets"],

"timestamp": datetime.utcnow().isoformat(),

},

{

"text": "New restaurant opens in downtown",

"category": "Food",

"tags": ["dining", "local"],

"timestamp": datetime.utcnow().isoformat(),

},

]

# Upload vectors with payloads

client.upsert(

collection_name=collection_name,

points=models.Batch(ids=list(range(num_vectors)), vectors=vectors, payloads=payloads),

)

print("Uploaded vectors with payloads")

# Search vectors

search_result = client.search(collection_name=collection_name, query_vector=vectors[0], limit=3)

print("\nSearch results:")

for scored_point in search_result:

print(

json.dumps(

{"id": scored_point.id, "score": scored_point.score, "payload": scored_point.payload}, indent=2

)

)

# Filtered search

filter_result = client.search(

collection_name=collection_name,

query_vector=vectors[0],

query_filter=models.Filter(

must=[models.FieldCondition(key="category", match=models.MatchValue(value="Technology"))]

),

limit=2,

)

print("\nFiltered search results:")

for scored_point in filter_result:

print(

json.dumps(

{"id": scored_point.id, "score": scored_point.score, "payload": scored_point.payload}, indent=2

)

)

# Create payload index

client.create_payload_index(

collection_name=collection_name, field_name="category", field_schema=models.PayloadFieldSchema.KEYWORD

)

print("\nCreated payload index on category field")

# Create vector index

client.create_payload_index(

collection_name=collection_name, field_name="tags", field_schema=models.PayloadFieldSchema.KEYWORD

)

print("Created payload index on tags field")

# Scroll through collection

scroll_result = client.scroll(collection_name=collection_name, limit=10, with_payload=True, with_vectors=True)

print("\nScrolled through collection:")

for point in scroll_result[0]:

print(json.dumps({"id": point.id, "payload": point.payload}, indent=2))

# Get collection info

collection_info = client.get_collection(collection_name)

print("\nCollection info:")

print(

json.dumps(

{

"name": collection_info.name,

"vectors_count": collection_info.vectors_count,

"points_count": collection_info.points_count,

"status": collection_info.status,

},

indent=2,

)

)

# Update payload

client.set_payload(collection_name=collection_name, payload={"new_field": "updated value"}, points=[0, 1])

print("\nUpdated payload for points 0 and 1")

# Delete points

client.delete(collection_name=collection_name, points_selector=models.PointIdsList(points=[4]))

print("Deleted point with id 4")

# Clean up

client.delete_collection(collection_name)

print("\nDeleted collection")

if __name__ == "__main__":

basic_example()

Back | FazBrowse Home | New Git URL