Milvus

pip install netra-sdk pymilvus

from netra import Netra
from pymilvus import MilvusClient
import os

# Initialize Netra
Netra.init(
    headers=f"x-api-key={os.environ.get('NETRA_API_KEY')}",
    trace_content=True
)

# Create Milvus client - automatically traced
client = MilvusClient(
    uri=os.environ.get('MILVUS_URI'),
    token=os.environ.get('MILVUS_TOKEN')
)

# Create collection
client.create_collection(
    collection_name="my_collection",
    dimension=384
)

from netra.decorators import task
from netra import SpanWrapper

@task()
def create_collection(client: MilvusClient, name: str, dimension: int):
    span = SpanWrapper("milvus-create-collection", {
        "collection.name": name,
        "vector.dimension": dimension
    }).start()
    
    client.create_collection(
        collection_name=name,
        dimension=dimension
    )
    
    span.end()

from netra.decorators import task
from netra import SpanWrapper, ActionModel

@task()
def insert_vectors(client: MilvusClient, collection: str, data: list):
    span = SpanWrapper("milvus-insert", {
        "collection": collection,
        "entities.count": len(data)
    }).start()
    
    result = client.insert(
        collection_name=collection,
        data=data
    )
    
    span.set_action([ActionModel(
        action="insert",
        action_type="database.insert",
        success=True,
        affected_records=[{"id": str(d["id"])} for d in data],
        metadata={"collection": collection, "inserted": result["insert_count"]}
    )])
    span.set_attribute("insert.count", result["insert_count"])
    span.end()
    
    return result

from netra.decorators import workflow
from netra import SpanWrapper

@workflow()
def search_vectors(client: MilvusClient, collection: str, query: list[float], limit: int = 5):
    span = SpanWrapper("milvus-search", {
        "collection": collection,
        "query.dimension": len(query),
        "limit": limit
    }).start()
    
    results = client.search(
        collection_name=collection,
        data=[query],
        limit=limit
    )
    
    span.set_attribute("results.count", len(results[0]))
    span.end()
    
    return results

from netra.decorators import task
from netra import SpanWrapper

@task()
def filter_search(client: MilvusClient, collection: str, query: list[float], filter: str):
    span = SpanWrapper("milvus-filter-search", {
        "collection": collection,
        "filter": filter
    }).start()
    
    results = client.search(
        collection_name=collection,
        data=[query],
        filter=filter,
        limit=10
    )
    
    span.set_attribute("results.count", len(results[0]))
    span.end()
    
    return results