ChromaDB

pip install netra-sdk chromadb

from netra import Netra
import chromadb
import os

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

# Create ChromaDB client - automatically traced
client = chromadb.Client()
collection = client.get_or_create_collection(name="my_collection")

# Add documents
collection.add(
    ids=["id1", "id2"],
    documents=["Document 1", "Document 2"],
    metadatas=[{"source": "web"}, {"source": "book"}]
)

from netra import task, SpanWrapper

@task()
def create_collection(client, name: str):
    span = SpanWrapper("chroma-create-collection", {
        "collection.name": name
    }).start()
    
    collection = client.create_collection(name=name)
    span.set_attribute("collection.id", collection.id)
    span.end()
    
    return collection

from netra import task, SpanWrapper, ActionModel

@task()
def add_documents(collection, documents: list[str]):
    span = SpanWrapper("chroma-add-documents", {
        "documents.count": len(documents)
    }).start()
    
    collection.add(
        ids=[f"doc{i}" for i in range(len(documents))],
        documents=documents
    )
    
    span.set_action([ActionModel(
        action="insert",
        action_type="database.insert",
        success=True,
        affected_records=[{"id": f"doc{i}"} for i in range(len(documents))],
        metadata={"collection": collection.name}
    )])
    span.set_attribute("status", "success")
    span.end()

from netra import workflow, SpanWrapper

@workflow()
def search_similar(collection, query: str, n: int = 5):
    span = SpanWrapper("chroma-search", {
        "query": query,
        "n_results": n
    }).start()
    
    results = collection.query(
        query_texts=[query],
        n_results=n
    )
    
    span.set_attribute("results.count", len(results["ids"][0]))
    span.end()
    
    return results

from netra import task, SpanWrapper
import json

@task()
def filter_search(collection, query: str, filter: dict):
    span = SpanWrapper("chroma-filter-search", {
        "query": query,
        "filter": json.dumps(filter)
    }).start()
    
    results = collection.query(
        query_texts=[query],
        where=filter
    )
    
    span.set_attribute("results.count", len(results["ids"][0]))
    span.end()
    
    return results