🚀 Feature: More Declarative Instrumentation Wrappers

[paolorechia]

Which component is this feature for?

All Packages

🔖 Feature description

We should improve our instrumentation wrappers to reduce the required amount of boilerplate per project.

The end-goal is to drive the instrumentations towards a more declarative approach, as much as possible. There will always be exceptions, so the solution also needs to be flexible with this.

🎤 Why is this feature needed ?

Less maintenance overhead.

✌️ How do you aim to achieve this?

We'll need to write some custom tooling for this.

In the weaviate instrumentation PR, I wrote some simple classes to reduce boilerplate, for instance:

class _BatchInstrumentor(_Instrumentor):
    namespace = "db.weaviate.batch"
    mapped_attributes = {
        "add_data_object": [
            "data_object",
            "class_name",
            "uuid",
            "vector",
            "tenant",
        ],
        "flush": [],
    }

We could try and take this further. Python offers an amazing introspection module in the standard library (inspect), so in theory just by passing the function name we should be able to automatically extract all arguments. That should be the default way to instrument methods.

Using such resources, we can then simply map in a declarative way most of the standard objects / functions, like we already do:

WRAPPED_METHODS = [
    {
        "package": chromadb.api.segment,
        "object": "SegmentAPI",
        "method": "_query",
        "span_name": "chroma.query.segment._query"
    },
    {
        "package": chromadb,
        "object": "Collection",
        "method": "add",
        "span_name": "chroma.add"
    },

But without the need to explicitly add calls to set span attributes, as our instrumentation tooling should automatically take care of that:

def _set_add_attributes(span, kwargs):
    _set_span_attribute(span, "db.chroma.add.ids_count", count_or_none(kwargs.get("ids")))
    _set_span_attribute(span, "db.chroma.add.embeddings_count", count_or_none(kwargs.get("embeddings")))
    _set_span_attribute(span, "db.chroma.add.metadatas_count", count_or_none(kwargs.get("metadatas")))
    _set_span_attribute(span, "db.chroma.add.documents_count", count_or_none(kwargs.get("documents")))

🔄️ Additional Information

No response

👀 Have you spent some time to check if this feature request has been raised before?

• I checked and didn't find similar issue

Are you willing to submit PR?

Yes I am willing to submit a PR!