Векторни бази данни

from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection
from pymilvus.model.dense import SentenceTransformerEmbeddingFunction

# 1. Свързване към Milvus (по подразбиране localhost:19530)
connections.connect("default", host="localhost", port="19530")

# 2. Инициализиране на функцията за ембединг с модела e5-base-v2
# Този модел изисква:
# - Префикс "passage: " за документите
# - Префикс "query: " за заявките за търсене
ef = SentenceTransformerEmbeddingFunction("intfloat/e5-base-v2")

# 3. Дефиниране на схемата на колекцията:
# - "id" — целочислен идентификатор (първичен ключ)
# - "text" — първоначалният текст на документа (низ)
# - "emb" — ембединг вектор с размерност 768
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=False),
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=512),
    FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=768)
]
schema = CollectionSchema(fields, description="Колекция с ембединги от e5-base-v2")

# 4. Създаване на колекцията в Milvus със зададената схема
collection = Collection("e5_collection", schema)

# 5. Подготовка и вмъкване на документи
# Важно: преди да се подадат на модела, трябва да се добави префиксът "passage: "
raw_docs = ["Hello world", "Milvus vector database", "Semantic search with e5 model"]
docs = [f"passage: {d}" for d in raw_docs]  # добавяме префикса
ids = [1, 2, 3]

# Изчисляваме ембедингите на документите с помощта на e5‑base
embs = ef.encode_documents(docs)

# Вмъкваме в колекцията:
# - идентификаторите
# - първоначалните (чисти) текстове без префикси
# - ембедингите
collection.insert([ids, raw_docs, embs])

# 6. Създаване на индекс по полето "emb" за ускоряване на търсенето
collection.create_index(
    field_name="emb",
    index_params={
        # тип на индекса
        "index_type": "IVF_FLAT",

        # параметър за разделяне на клъстери
        "params": {"nlist": 128},

        # метрика за разстоянието (евклидово разстояние)
        "metric_type": "L2"
    }
)

# 7. Зареждане на колекцията в оперативната памет
# Без това търсенето няма да работи
collection.load()

# 8. Заявка за търсене
# По същия начин – използваме префикса "query: " пред текста на заявката
query_docs = ["query: vector database"]
q_emb = ef.encode_queries(query_docs)

# Изпълняваме семантично търсене по ембедингите
results = collection.search(
    # ембединг на заявката за търсене
    data=q_emb,

    # поле, по което се извършва търсенето
    anns_field="emb",

    # параметри на търсенето
    param={"metric_type": "L2", "params": {"nprobe": 10}},

    # брой най-близки съседи
    limit=2,

    # допълнителни полета, които трябва да се върнат
    output_fields=["text"]
)

# 9. Извеждаме резултатите от търсенето
for i, hits in enumerate(results):
    print(f"Резултати за заявката: '{query_docs[i]}'")
    if not hits:
        print("Нищо не е намерено")
        continue
    for rank, hit in enumerate(hits, start=1):
        print(f"  {rank}:")
        print(f"    ID: {hit.id}")
        print(f"    Текст: {hit.entity.get('text')}")
        print(f"    Разстояние: {hit.distance:.4f}")

# Резултати за заявката: 'query: vector database'
#  1:
#    ID: 2
#    Текст: Milvus vector database
#    Разстояние: 2.8374
#  2:
#    ID: 3
#    Текст: Semantic search with e5 model
#    Разстояние: 5.4931

# 10. Изтриване на документ по ID
# В случая се изтрива документът с id = 1
collection.delete(expr="id in [1]")

# 11. Изтриване на цялата колекция (ако вече не е нужна)
collection.drop()