Milvus

Exemple de collecte unique

Créez un fichier nommé 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()

Exemple de collections multiples

Cet exemple utilise pymilvus trois collections indépendantes, chacune s'appuyant sur une PDF différente. Il couvre create_collection, create_index, insert, search, query,, upsert, deletedescribe_collection, et list_collections. Il intègre également une fonctionnalité de recherche inter-collections qui regroupe les résultats en une seule réponse générée par le modèle LLM « watsonx ». Traceloop @task et les décorateurs @workflow enregistrent chaque opération effectuée dans Instana sous la forme d'un « span » nommé.

Créez un fichier nommé 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()

Configuration des identifiants d' IBM watsonx

Pour accéder à IBM watsonx, définissez les variables d'environnement suivantes :

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

Exécution de l'application

Pour exécuter l'exemple de collecte unique, exécutez la commande suivante :

python WatsonxEmbeddingMilvus.py

Pour exécuter l'exemple RAG avec plusieurs collections, exécutez la commande suivante :

python rag-watsonx-traceloop.py

Exportation des données de télémétrie vers Instana

Pour exporter des traces et des métriques vers Instana, utilisez la configuration d'exportation suivante :