File size: 1,500 Bytes
af9408a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import datetime
import subprocess
from qdrant_client import QdrantClient, models
import os
class QdrantClientInitializer:

    def __init__(self):
        pass

    def run_docker(self):
        # Define the Docker command
        docker_command = [
            "docker run -d -p 6333:6333 -p 6334:6334 \
                -v $(pwd)/qdrant_storage:/qdrant/storage:z \
                    qdrant/qdrant"
        ]

        # Run the Docker command
        subprocess.run(docker_command, shell=True)
        print("Docker container is running in the background.")

    def initialize_db(self):
        client = QdrantClient(
            url=os.getenv('qdrant_url'), 
            api_key=os.getenv('qdrant_api'), timeout=20)
        # client = QdrantClient(url="http://localhost:6333", timeout=1200, grpc_port=6333)

        return client

    def create_collection(self, client, collection_name, vector_size=4096, distance=models.Distance.COSINE):

        if client.collection_exists(collection_name=collection_name):
            print("{name} is exist!".format(name=collection_name))

        else:
            print("Name is not exist!")
            client.create_collection(collection_name=collection_name,
                                     vectors_config=models.VectorParams(size=vector_size, 
                                                                        distance=distance,
                                                                        )
                                    )