testbot_v4 / initialize_db.py
soyleyicicem's picture
Upload 8 files
af9408a verified
raw
history blame
1.5 kB
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,
)
)