Vector databases

[0.12, -0.56, 0.44, ..., -0.03]

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

# 1. Connect to Milvus (default localhost:19530)
connections.connect("default", host="localhost", port="19530")

#2. Initialize the embedding function with the e5-base-v2 model
# This model requires:
# - Prefix "passage: " for documents
# - Prefix "query: " for search queries
ef = SentenceTransformerEmbeddingFunction("intfloat/e5-base-v2")

# 3. Defining the collection schema:
# - "id" — integer identifier (primary key)
# - "text" — the original text of the document (string)
# - "emb" — embedding vector of dimension 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="Collection with embeddings from e5-base-v2")

# 4. Creating a collection in Milvus with a given schema
collection = Collection("e5_collection", schema)

# 5. Preparing and inserting documents
# Important: Before submitting to the model, you need to add the prefix "passage: "
raw_docs = ["Hello world", "Milvus vector database", "Semantic search with e5 model"]
docs = [f"passage: {d}" for d in raw_docs] # add prefix
ids = [1, 2, 3]

# Calculate embeddings for documents using e5‑base
embs = ef.encode_documents(docs)

# Insert into the collection:
# - identifiers
# - original (clean) texts without prefixes
# - embeddings
collection.insert([ids, raw_docs, embs])

# 6. Create an index on the "emb" field to speed up searching
collection.create_index(
    field_name="emb",
    index_params={
        # index type
        "index_type": "IVF_FLAT",

        # clustering parameter
        "params": {"nlist": 128},

        # distance metric (euclidean distance)
        "metric_type": "L2"
    }
)

# 7. Loading the collection into RAM
# Without this the search will not work
collection.load()

# 8. Search query
# Similarly, we use the prefix "query: " before the query text
query_docs = ["query: vector database"]
q_emb = ef.encode_queries(query_docs)

# Performing semantic search by embeddings
results = collection.search(
    # search query embedding
    data=q_emb,

    # field by which the search is performed
    anns_field="emb",

    # search parameters
    param={"metric_type": "L2", "params": {"nprobe": 10}},

    # number of nearest neighbors
    limit=2,

    # additional fields to return
    output_fields=["text"]
)

# 9. Display search results
for i, hits in enumerate(results):
    print(f"Results for query: '{query_docs[i]}'")
    if not hits:
        print("Nothing found")
        continue
    for rank, hit in enumerate(hits, start=1):
        print(f"  {rank}:")
        print(f"    ID: {hit.id}")
        print(f" Text: {hit.entity.get('text')}")
        print(f" Distance: {hit.distance:.4f}")

# Results for query: 'query: vector database'
#  1:
#    ID: 2
# Text: Milvus vector database
# Distance: 2.8374
#  2:
#    ID: 3
# Text: Semantic search with e5 model
# Distance: 5.4931

# 10. Deleting a document by ID
# In this case, the document with id = 1 is deleted
collection.delete(expr="id in [1]")

# 11. Delete the entire collection (if no longer needed)
collection.drop()