taspol commited on
Commit
8ffea85
·
1 Parent(s): faaabaa

fix: Restqdrant

Browse files
Files changed (1) hide show
  1. class_mod/rest_qdrant.py +39 -0
class_mod/rest_qdrant.py CHANGED
@@ -15,6 +15,45 @@ class RestQdrantClient:
15
  r.raise_for_status()
16
  return r.json()
17
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18
  def upsert(self, collection_name, points):
19
  r = self.session.put(
20
  f"{self.url}/collections/{collection_name}/points",
 
15
  r.raise_for_status()
16
  return r.json()
17
 
18
+ def get_collection(self, collection_name):
19
+ r = self.session.get(f"{self.url}/collections/{collection_name}", timeout=self.timeout)
20
+ r.raise_for_status()
21
+ return r.json()
22
+ def search(self, collection_name, query_vector, limit=10, with_payload=True):
23
+ payload = {
24
+ "vector": query_vector,
25
+ "limit": limit,
26
+ "with_payload": with_payload
27
+ }
28
+ r = self.session.post(
29
+ f"{self.url}/collections/{collection_name}/points/search",
30
+ json=payload,
31
+ timeout=self.timeout
32
+ )
33
+ r.raise_for_status()
34
+ return r.json()
35
+
36
+ def delete_collection(self, collection_name):
37
+ r = self.session.delete(f"{self.url}/collections/{collection_name}", timeout=self.timeout)
38
+ if r.status_code not in [200, 404]: # 404 means collection didn't exist
39
+ r.raise_for_status()
40
+ return r.json() if r.text else {}
41
+
42
+ def create_collection(self, collection_name, vector_size, distance="Cosine"):
43
+ payload = {
44
+ "vectors": {
45
+ "size": vector_size,
46
+ "distance": distance.upper() # "COSINE", "EUCLIDEAN", "DOT"
47
+ }
48
+ }
49
+ r = self.session.put(f"{self.url}/collections/{collection_name}", json=payload, timeout=self.timeout)
50
+ r.raise_for_status()
51
+ return r.json()
52
+ def recreate_collection(self, collection_name, vector_size, distance="Cosine"):
53
+ # Delete if exists
54
+ self.delete_collection(collection_name)
55
+ # Create new collection
56
+ return self.create_collection(collection_name, vector_size, distance)
57
  def upsert(self, collection_name, points):
58
  r = self.session.put(
59
  f"{self.url}/collections/{collection_name}/points",