37 lines
1.3 KiB
Python
37 lines
1.3 KiB
Python
import os
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.models import Distance, VectorParams, PointStruct
|
|
|
|
class VectorDB:
|
|
def __init__(self, collection_name="wifi_lab", vector_size=4096):
|
|
host = os.getenv("QDRANT_HOST", "localhost")
|
|
port = int(os.getenv("QDRANT_PORT", 6333))
|
|
self.client = QdrantClient(host=host, port=port)
|
|
self.collection_name = collection_name
|
|
|
|
# Создаем коллекцию, если её нет (размерность 4096 для Qwen 8B Embedding)
|
|
self.client.recreate_collection(
|
|
collection_name=self.collection_name,
|
|
vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE),
|
|
)
|
|
|
|
def add_packet(self, pkt_id, vector, metadata, text):
|
|
self.client.upsert(
|
|
collection_name=self.collection_name,
|
|
points=[
|
|
PointStruct(
|
|
id=pkt_id,
|
|
vector=vector,
|
|
payload={"metadata": metadata, "text": text}
|
|
)
|
|
]
|
|
)
|
|
|
|
def find_similar(self, vector):
|
|
return self.client.search(
|
|
collection_name=self.collection_name,
|
|
query_vector=vector,
|
|
limit=3,
|
|
with_payload=True
|
|
)
|