Milvus

Single-collection example

Create a file named WatsonxEmbeddingMilvus.py:

from pymilvus import MilvusClient

from traceloop.sdk import Traceloop
from traceloop.sdk.decorators import workflow, task

from langchain_ibm.embeddings import WatsonxEmbeddings 
from ibm_watsonx_ai.metanames import EmbedTextParamsMetaNames
import os

Traceloop.init(app_name="Watsonx_Embeddings_MilvusClient")

# connect to Milvus Locally
@task(name="setup_milvus_client")
def setup_milvus_client(uri: str, collection_name: str, dimension: int):
    client = MilvusClient(uri=uri)
    ##create a collection in Milvus DB
    if client.has_collection(collection_name=collection_name):
        client.drop_collection(collection_name=collection_name)
    client.create_collection(
        collection_name=collection_name, dimension=dimension, timeout=10, metric_type="COSINE"
    )
    return client


embedding_model = None  # Define Embedding Model globally

#Initialize watsonx embedding model 
@task(name="initialize_embedding_model")
def initialize_embedding_model(
    ibm_cloud_url: str,
    ibm_cloud_api_key: str,
    model_id: str,
    project_id: str,
    model_kwargs: dict = None,
    encode_kwargs: dict = None,
):
    embed_params = {
    EmbedTextParamsMetaNames.TRUNCATE_INPUT_TOKENS: 3,
    EmbedTextParamsMetaNames.RETURN_OPTIONS: {"input_text": True},
    }
    global embedding_model
    model_kwargs = model_kwargs or {}
    encode_kwargs = encode_kwargs or {"normalize_embeddings": False}

    embedding_model = WatsonxEmbeddings(
        url=ibm_cloud_url,
        project_id=project_id,
        model_id=model_id,
        apikey=ibm_cloud_api_key,
        params=embed_params
    )


# embed Documents and insert it to Milvus DB
@task(name="encode_documents_and_insert")
def encode_documents_and_insert(
    client: MilvusClient,
    collection_name: str,
    partition_name: str,
    docs: list,
    subject: str,
    timeout: float,
):
    vectors = embedding_model.embed_documents(docs)
    data = [
        {"id": i, "vector": vectors[i], "text": docs[i], "subject": subject}
        for i in range(len(vectors))
    ]

    res = client.insert(
        collection_name=collection_name,
        partition_name=partition_name,
        data=data,
        timeout=timeout,
    )
    print(res)


# apply vector embedding on the query and search the same in the vecotr db
@task(name="perform_vector_search")
def perform_vector_search(
    client: MilvusClient,
    collection_name: str,
    query: str,
    limit: int,
    output_fields: list,
):
    query_vector = embedding_model.embed_query(query)
    result = client.search(
        collection_name=collection_name,
        partition_name="partitionA",
        data=[query_vector],
        limit=limit,
        output_fields=output_fields,
    )
    return result


# search in vecotr db with filters applied
@task(name="perform_vector_search_with_filter")
def perform_vector_search_with_filter(
    client: MilvusClient,
    collection_name: str,
    partition_names: list,
    anns_field: str,
    search_params: dict,
    query: str,
    filter: str,
    limit: int,
    output_fields: list,
    timeout: float,
):
    query_vector = embedding_model.embed_query(query)
    searchResult = client.search(
        collection_name=collection_name,
        partition_names=partition_names,
        search_params=search_params,
        anns_field=anns_field,
        data=[query_vector],
        filter=filter,
        limit=limit,
        output_fields=output_fields,
        timeout=timeout,
    )
    return searchResult


# query for entries in the Collection 
@task(name="perform_query")
def perform_query(
    client: MilvusClient,
    collection_name: str,
    filter: str,
    output_fields: list,
):
    queryResult = client.query(
        collection_name=collection_name,
        filter=filter,
        partition_names=["partitionA"],
        output_fields=output_fields,
    )
    return queryResult


# query the db passing list of ids
@task(name="perform_query_ids")
def perform_query_Ids_partition(
    client: MilvusClient,
    collection_name: str,
    partition_names: list,
    limit: int,
    ids: list,
    output_fields: list,
    timeout: float,
):
    queryResult = client.query(
        collection_name=collection_name,
        partition_names=partition_names,
        limit=limit,
        ids=ids,
        timeout=timeout,
    )
    return queryResult

# delete entries from the collection
@task(name="delete_entities")
def delete_entities(
    client: MilvusClient,
    collection_name: str,
    partition_name: str,
    ids: list = None,
    filter: str = None,
    timeout: float = None,
):
    if ids is not None:
        deleteResult = client.delete(collection_name=collection_name, ids=ids)
        print(deleteResult)
    if filter is not None:
        deleteRes = client.delete(
            collection_name=collection_name,
            timeout=timeout,
            filter=filter,
            partition_name=partition_name,
        )
        print(deleteRes)

# modify data in the collection
@task(name="upsert_entities")
def upsert_entities( 
    client: MilvusClient,
    collection_name: str,
    partition_name: str,
    docs: list,
    ids: list,
    subject: str,
    timeout: float,
):
    vectors = embedding_model.embed_documents(docs)
    data = [
        {"id": ids[i], "vector": vectors[i], "text": docs[i], "subject": subject}
        for i in range(len(vectors))
    ]

    res = client.upsert( 
        collection_name=collection_name,
        partition_name=partition_name,
        data=data,
        timeout=timeout,
    )
    print("Upsert Result:", res)


@task(name="get_entities")
def get_entities(
    client: MilvusClient,
    collection_name: str,
    partition_names: list,
    output_fields: list,
    ids: list,
    timeout: float,
):
    result = client.get(
        collection_name=collection_name,
        partition_names=partition_names,
        output_fields=output_fields,
        ids=ids,
        timeout=timeout,
    )
    return result


@workflow(name="milvus_operations_with_watsonx")  
def milvus_operations_with_watsonx():
    client = setup_milvus_client(
        uri="http://127.0.0.1:19530", collection_name="demo_collection", dimension=768
    )
    partition_name = "partitionA"
    client.create_partition(
        collection_name="demo_collection", partition_name=partition_name
    )

    #  Watsonx Embedding model parameters
    ibm_cloud_url = os.getenv("WATSONX_URL")
    ibm_cloud_api_key = os.getenv("WATSONX_API_KEY")
    model_id = (
        "ibm/slate-125m-english-rtrvr"  # or any other supported model
    )
    project_id=os.getenv("WATSONX_PROJECT_ID")

    initialize_embedding_model(
        ibm_cloud_url=ibm_cloud_url,
        ibm_cloud_api_key=ibm_cloud_api_key,
        model_id=model_id,
        project_id=project_id
    )

    docs_history = [
        "Artificial intelligence was founded as an academic discipline in 1956.",
        "Alan Turing was the first person to conduct substantial research in AI.",
        "Born in Maida Vale, London, Turing was raised in southern England.",
    ]

    encode_documents_and_insert(
        client=client,
        collection_name="demo_collection",
        partition_name=partition_name,
        docs=docs_history,
        subject="history",
        timeout=10,
    )  

    # Upsert example
    new_docs_history = [
        "Alan Turing developed the Turing Test.",
        "Artificial intelligence continues to evolve.",
    ]
    new_ids_history = [
        0,
        1,
    ]  
    upsert_entities(
        client=client,
        collection_name="demo_collection",
        partition_name=partition_name,
        docs=new_docs_history,
        ids=new_ids_history,
        subject="history",
        timeout=10,
    )

    # Get example
    get_result = get_entities(
        client=client,
        collection_name="demo_collection",
        partition_names=[partition_name],
        output_fields=["text", "subject"],
        ids=new_ids_history,
        timeout=10,
    )
    print("Get Result:", get_result)

    # Semantic Search
    # Vector search
    result = perform_vector_search(
        client=client,
        collection_name="demo_collection",
        query="Who is Alan Turing?",
        limit=2,
        output_fields=["text", "subject"],
    )
    print(result)

    # Vector Search with Metadata Filtering
    docs_biology = [
        "Machine learning has been used for drug design.",
        "Computational synthesis with AI algorithms predicts molecular properties.",
        "DDR1 is involved in cancers and fibrosis.",
    ]

    encode_documents_and_insert(
        client=client,
        collection_name="demo_collection",
        partition_name=partition_name,
        docs=docs_biology,
        subject="biology",
        timeout=10,
    )

    search_params = {"metric_type": "COSINE", "params": {}}

    searchResult = perform_vector_search_with_filter(
        client=client,
        collection_name="demo_collection",
        partition_names=[partition_name],
        anns_field="vector",
        search_params=search_params,
        query="tell me AI related information",
        filter="subject == 'biology'",
        limit=2,
        output_fields=["text", "subject"],
        timeout=10,
    )
    print(searchResult)

    # Perform Query
    queryResult = perform_query(
        client=client,
        collection_name="demo_collection",
        filter="subject == 'history'",
        output_fields=["text", "subject"],
    )
    print(queryResult)

    # Perform Query with ids as input param
    queryResult = perform_query_Ids_partition(
        client=client,
        collection_name="demo_collection",
        partition_names=[partition_name],
        limit=1,
        ids=[0, 2],
        output_fields=["text", "subject"],
        timeout=10,
    )
    print(queryResult)

    # Delete entities
    delete_entities(
        client=client,
        collection_name="demo_collection",
        partition_name=partition_name,
        ids=[0, 2],
        timeout=10,
    )

    # 8. Delete entities by a filter expression
    delete_entities(
        client=client,
        collection_name="demo_collection",
        partition_name=partition_name,
        filter="subject == 'biology'",
        timeout=10,
    )


milvus_operations_with_watsonx()

Multiple-collections example

This example uses pymilvus with three independent collections, each backed by a different PDF. It covers create_collection, create_index, insert, search, query, upsert, delete, describe_collection, and list_collections. It also includes a cross-collection search that merges results into a single watsonx LLM response. Traceloop @task and @workflow decorators trace every operation in Instana as a named span.

Create a file named rag-watsonx-traceloop.py:

"""RAG with WatsonX + PyMilvus (multiple collections), instrumented by Traceloop."""

import os
import warnings

from ibm_watsonx_ai.wml_resource import WatsonxAPIWarning
warnings.filterwarnings("ignore", category=WatsonxAPIWarning)

from langchain_ibm import WatsonxLLM
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from pymilvus import MilvusClient, DataType
from traceloop.sdk import Traceloop
from traceloop.sdk.decorators import task, workflow

Traceloop.init(app_name="rag_pymilvus_watsonx_service")

os.environ["WATSONX_URL"]        = "<watsonx-url>"
os.environ["WATSONX_API_KEY"]    = "<watsonx-api-key>"
os.environ["WATSONX_PROJECT_ID"] = "<watsonx-project-id>"

MILVUS_URI          = f"http://{os.getenv('MILVUS_HOST', '127.0.0.1')}:{os.getenv('MILVUS_PORT', '19530')}"
EMBEDDING_DIM       = 768
UPSERT_SYNTHETIC_ID = 999999

COLLECTIONS = [
    {
        "name": "milvus_docs_collection",
        "pdf":  os.path.expanduser("~/Downloads/milvus-overview.pdf"),
        "questions": [
            "What are the main features of Milvus?",
            "Summarize the architecture overview.",
        ],
    },
    {
        "name": "ibm_docs_collection",
        "pdf":  os.path.expanduser("~/Downloads/ibm-report.pdf"),
        "questions": [
            "What are the key topics covered in this document?",
            "Summarize the main findings or recommendations.",
        ],
    },
    {
        "name": "otel_docs_collection",
        "pdf":  os.path.expanduser("~/Downloads/opentelemetry-overview.pdf"),
        "questions": [
            "What is OpenTelemetry and what does this document describe?",
            "What are the key instrumentation concepts mentioned?",
        ],
    },
]


@task(name="initialize_watsonx_llm")
def initialize_llm() -> WatsonxLLM:
    return WatsonxLLM(
        model_id="meta-llama/llama-3-3-70b-instruct",
        url=os.getenv("WATSONX_URL"),
        apikey=os.getenv("WATSONX_API_KEY"),
        project_id=os.getenv("WATSONX_PROJECT_ID"),
        params={"max_new_tokens": 512, "temperature": 0.5, "top_k": 50, "top_p": 1},
    )


@task(name="create_collection")
def create_collection(client: MilvusClient, collection_name: str):
    if client.has_collection(collection_name):
        client.drop_collection(collection_name)
    schema = client.create_schema(auto_id=False, enable_dynamic_field=True)
    schema.add_field("id",        DataType.INT64,        is_primary=True, auto_id=False)
    schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=EMBEDDING_DIM)
    schema.add_field("text",      DataType.VARCHAR,      max_length=4096)
    schema.add_field("source",    DataType.VARCHAR,      max_length=512)
    schema.add_field("page",      DataType.INT64)
    client.create_collection(collection_name=collection_name, schema=schema)


@task(name="create_index")
def create_index(client: MilvusClient, collection_name: str):
    index_params = client.prepare_index_params()
    index_params.add_index(
        field_name="embedding",
        index_type="HNSW",
        metric_type="COSINE",
        params={"M": 16, "efConstruction": 200},
    )
    client.create_index(collection_name=collection_name, index_params=index_params)
    client.load_collection(collection_name=collection_name)


@task(name="process_and_insert_documents")
def process_and_insert(client, collection_name, pdf_path, embeddings) -> int:
    chunks = RecursiveCharacterTextSplitter(
        chunk_size=1000, chunk_overlap=100
    ).split_documents(PyPDFLoader(pdf_path).load())
    rows = [
        {
            "id":        idx,
            "embedding": embeddings.embed_query(c.page_content),
            "text":      c.page_content,
            "source":    c.metadata.get("source", pdf_path),
            "page":      int(c.metadata.get("page", 0)),
        }
        for idx, c in enumerate(chunks)
    ]
    result = client.insert(collection_name=collection_name, data=rows)
    return result["insert_count"]


@task(name="describe_collection")
def describe_collection(client: MilvusClient, collection_name: str):
    info  = client.describe_collection(collection_name)
    stats = client.get_collection_stats(collection_name)
    print(f"  fields: {[f['name'] for f in info['fields']]}, rows: {stats['row_count']}")
    return stats


@task(name="vector_search")
def vector_search(client, collection_name, query, embeddings, top_k=3) -> list:
    results = client.search(
        collection_name=collection_name,
        data=[embeddings.embed_query(query)],
        limit=top_k,
        output_fields=["text", "page", "source"],
        search_params={"metric_type": "COSINE", "params": {"ef": 64}},
    )
    return [hit["entity"]["text"] for hit in results[0]]


@task(name="scalar_query")
def scalar_query(client: MilvusClient, collection_name: str, page: int = 0) -> list:
    return client.query(
        collection_name=collection_name,
        filter=f"page == {page}",
        output_fields=["id", "text", "page"],
        limit=5,
    )


@task(name="upsert_chunk")
def upsert_chunk(client, collection_name, embeddings):
    text = "Updated synthetic chunk for upsert demonstration."
    client.upsert(
        collection_name=collection_name,
        data=[{
            "id":        UPSERT_SYNTHETIC_ID,
            "embedding": embeddings.embed_query(text),
            "text":      text,
            "source":    "synthetic",
            "page":      9999,
        }],
    )


@task(name="delete_by_filter")
def delete_by_filter(client: MilvusClient, collection_name: str):
    client.delete(collection_name=collection_name, filter=f"id == {UPSERT_SYNTHETIC_ID}")


@task(name="rag_answer")
def rag_answer(llm: WatsonxLLM, context_chunks: list, question: str) -> str:
    context = "\n\n".join(context_chunks)
    prompt  = f"Context:\n{context}\n\nQuestion: {question}\n\nAnswer:"
    answer  = llm.invoke(prompt)
    print(f"  Q: {question}\n  A: {answer}")
    return answer


@task(name="cross_collection_search")
def cross_collection_search(client, embeddings, llm, question):
    all_chunks = []
    for cfg in COLLECTIONS:
        all_chunks.extend(vector_search(client, cfg["name"], question, embeddings, top_k=2))
    rag_answer(llm, all_chunks, question)


@workflow(name="rag_pymilvus_multi_collection_workflow")
def rag_workflow():
    llm        = initialize_llm()
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
    client     = MilvusClient(uri=MILVUS_URI)

    for cfg in COLLECTIONS:
        name = cfg["name"]
        create_collection(client, name)
        create_index(client, name)
        process_and_insert(client, name, cfg["pdf"], embeddings)
        describe_collection(client, name)
        for question in cfg["questions"]:
            rag_answer(llm, vector_search(client, name, question, embeddings), question)
        scalar_query(client, name, page=0)
        upsert_chunk(client, name, embeddings)
        delete_by_filter(client, name)

    print(f"All collections: {client.list_collections()}")
    cross_collection_search(
        client, embeddings, llm,
        question="Give a brief summary of the most important information.",
    )


if __name__ == "__main__":
    rag_workflow()

Configuring IBM watsonx credentials

To access IBM watsonx, set the following environment variables:

export WATSONX_URL=<watsonx-url>
export WATSONX_API_KEY=<watsonx-iam-api-key>
export WATSONX_PROJECT_ID=<watsonx-project-id>

Running the application

To run the single-collection example, run the following command:

python WatsonxEmbeddingMilvus.py

To run the multiple-collections RAG example, run the following command:

python rag-watsonx-traceloop.py

Exporting telemetry to Instana

To export traces and metrics to Instana, use the following export configuration: