fix(milvus):Updated the instrumentation to collect get() and create_collection() span attributes

packages/opentelemetry-instrumentation-milvus/opentelemetry/instrumentation/milvus/__init__.py

@@ -20,6 +20,12 @@
 _instruments = ("pymilvus >= 2.4.1",)
 
 WRAPPED_METHODS = [
+    {
+        "package": pymilvus,
+        "object": "MilvusClient",
+        "method": "create_collection",
+        "span_name": "milvus.create_collection"
+    },
     {
         "package": pymilvus,
         "object": "MilvusClient",

packages/opentelemetry-instrumentation-milvus/opentelemetry/instrumentation/milvus/wrapper.py

@@ -51,6 +51,8 @@ def _wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
             _set_get_attributes(span, kwargs)
         elif to_wrap.get("method") == "query":
             _set_query_attributes(span, kwargs)
+        elif to_wrap.get("method") == "create_collection":
+            _set_create_collection_attributes(span, kwargs)
 
         return_value = wrapped(*args, **kwargs)
         if to_wrap.get("method") == "query":
@@ -89,7 +91,45 @@ def count_or_none(obj):
 
     return None
 
-
+@dont_throw
+def  _set_create_collection_attributes(span, kwargs):
+    _set_span_attribute(
+        span,
+        AISpanAttributes.MILVUS_CREATE_COLLECTION_NAME,
+        kwargs.get("collection_name"),
+    )
+    _set_span_attribute(
+        span,
+        AISpanAttributes.MILVUS_CREATE_COLLECTION_DIMENSION,
+        kwargs.get("dimension"),
+    )
+    _set_span_attribute(
+        span,
+        AISpanAttributes.MILVUS_CREATE_COLLECTION_PRIMARY_FIELD,
+        kwargs.get("primary_field_name"),
+    )
+    _set_span_attribute(
+        span,
+        AISpanAttributes.MILVUS_CREATE_COLLECTION_METRIC_TYPE,
+        kwargs.get("metric_type"),
+    )
+    _set_span_attribute(
+        span,
+        AISpanAttributes.MILVUS_CREATE_COLLECTION_TIMEOUT,
+        kwargs.get("timeout"),
+    )
+    _set_span_attribute(
+        span,
+        AISpanAttributes.MILVUS_CREATE_COLLECTION_ID_TYPE,
+        kwargs.get("id_type"),
+    )
+    _set_span_attribute(
+        span,
+        AISpanAttributes.MILVUS_CREATE_COLLECTION_VECTOR_FIELD,
+        kwargs.get("vector_field_name"),
+    )
+
+
 @dont_throw
 def _set_insert_attributes(span, kwargs):
     _set_span_attribute(
@@ -118,15 +158,15 @@ def _set_get_attributes(span, kwargs):
         span, AISpanAttributes.MILVUS_GET_COLLECTION_NAME, kwargs.get("collection_name")
     )
     _set_span_attribute(
-        span, AISpanAttributes.MILVUS_QUERY_IDS_COUNT, count_or_none(kwargs.get("ids"))
+        span, AISpanAttributes.MILVUS_GET_IDS_COUNT, count_or_none(kwargs.get("ids"))
     )
     _set_span_attribute(
         span,
-        AISpanAttributes.MILVUS_SEARCH_OUTPUT_FIELDS_COUNT,
+        AISpanAttributes.MILVUS_GET_OUTPUT_FIELDS_COUNT,
         count_or_none(kwargs.get("output_fields")),
     )
     _set_span_attribute(
-        span, AISpanAttributes.MILVUS_INSERT_TIMEOUT, kwargs.get("timeout")
+        span, AISpanAttributes.MILVUS_GET_TIMEOUT, kwargs.get("timeout")
     )
     _set_span_attribute(
         span,
@@ -165,7 +205,7 @@ def _set_search_attributes(span, kwargs):
     _set_span_attribute(
         span,
         AISpanAttributes.MILVUS_SEARCH_PARTITION_NAMES_COUNT,
-        kwargs.get("partition_name"),
+        count_or_none(kwargs.get("partition_names")),
     )
     _set_span_attribute(
         span, AISpanAttributes.MILVUS_SEARCH_ANNS_FIELD, kwargs.get("anns_field")
@@ -221,8 +261,8 @@ def _set_upsert_attributes(span, kwargs):
     )
     _set_span_attribute(
         span,
-        AISpanAttributes.MILVUS_UPSERT_TIMEOUT_COUNT,
-        count_or_none(kwargs.get("timeout")),
+        AISpanAttributes.MILVUS_UPSERT_TIMEOUT,
+        kwargs.get("timeout"),
     )
     _set_span_attribute(
         span,
@@ -240,8 +280,8 @@ def _set_delete_attributes(span, kwargs):
     )
     _set_span_attribute(
         span,
-        AISpanAttributes.MILVUS_DELETE_TIMEOUT_COUNT,
-        count_or_none(kwargs.get("timeout")),
+        AISpanAttributes.MILVUS_DELETE_TIMEOUT,
+        kwargs.get("timeout"),
     )
     _set_span_attribute(
         span,

packages/opentelemetry-instrumentation-milvus/tests/test_query.py

@@ -196,3 +196,33 @@ def test_milvus_query_and(exporter, collection):
         _id = event.attributes.get("id")
         assert isinstance(tag, str)
         assert isinstance(_id, int)
+
+
+def test_milvus_create_collection(exporter, collection):
+    collection_name = "quick_setup"
+    milvus.create_collection(
+        collection_name=collection_name,
+        dimension=5,
+        primary_field_name="my_id",
+        id_type="string",
+        vector_field_name="my_vector",
+        metric_type="L2",
+        auto_id=True,
+        max_length=512,
+        timeout=10
+    )
+    spans = exporter.get_finished_spans()
+    span = next(span for span in spans if span.name == "milvus.create_collection")
+
+    assert span.attributes.get(SpanAttributes.VECTOR_DB_VENDOR) == "milvus"
+    assert span.attributes.get(SpanAttributes.VECTOR_DB_OPERATION) == "create_collection"
+    assert (
+        span.attributes.get(SpanAttributes.MILVUS_CREATE_COLLECTION_NAME) == collection_name
+    )
+    assert span.attributes.get(SpanAttributes.MILVUS_CREATE_COLLECTION_DIMENSION) == 5
+    assert span.attributes.get(SpanAttributes.MILVUS_CREATE_COLLECTION_ID_TYPE) == "string"
+    assert span.attributes.get(SpanAttributes.MILVUS_CREATE_COLLECTION_METRIC_TYPE) == "L2"
+    assert span.attributes.get(SpanAttributes.MILVUS_CREATE_COLLECTION_TIMEOUT) == 10
+    assert span.attributes.get(SpanAttributes.MILVUS_CREATE_COLLECTION_PRIMARY_FIELD) == "my_id"
+    assert span.attributes.get(SpanAttributes.MILVUS_CREATE_COLLECTION_VECTOR_FIELD) == "my_vector"
+

packages/opentelemetry-semantic-conventions-ai/opentelemetry/semconv_ai/__init__.py

@@ -162,9 +162,19 @@ class SpanAttributes:
     MILVUS_DELETE_FILTER = "db.milvus.delete.filter"
     MILVUS_DELETE_IDS_COUNT = "db.milvus.delete.ids_count"
     MILVUS_DELETE_PARTITION_NAME = "db.milvus.delete.partition_name"
-    MILVUS_DELETE_TIMEOUT_COUNT = "db.milvus.delete.timeout_count"
+    MILVUS_DELETE_TIMEOUT = "db.milvus.delete.timeout"
     MILVUS_GET_COLLECTION_NAME = "db.milvus.get.collection_name"
     MILVUS_GET_PARTITION_NAMES_COUNT = "db.milvus.get.partition_names_count"
+    MILVUS_GET_IDS_COUNT = "db.milvus.get.ids_count"
+    MILVUS_GET_OUTPUT_FIELDS_COUNT = "db.milvus.get.output_fields_count"    
+    MILVUS_GET_TIMEOUT = "db.milvus.get.timeout"
+    MILVUS_CREATE_COLLECTION_NAME = "db.milvus.create_collection.collection_name"
+    MILVUS_CREATE_COLLECTION_DIMENSION = "db.milvus.create_collection.dimension"
+    MILVUS_CREATE_COLLECTION_PRIMARY_FIELD = "db.milvus.create_collection.primary_field"
+    MILVUS_CREATE_COLLECTION_METRIC_TYPE = "db.milvus.create_collection.metric_type"
+    MILVUS_CREATE_COLLECTION_TIMEOUT = "db.milvus.create_collection.timeout"
+    MILVUS_CREATE_COLLECTION_ID_TYPE = "db.milvus.create_collection.id_type"
+    MILVUS_CREATE_COLLECTION_VECTOR_FIELD = "db.milvus.create_collection.vector_field"
     MILVUS_INSERT_COLLECTION_NAME = "db.milvus.insert.collection_name"
     MILVUS_INSERT_DATA_COUNT = "db.milvus.insert.data_count"
     MILVUS_INSERT_PARTITION_NAME = "db.milvus.insert.partition_name"
@@ -188,7 +198,7 @@ class SpanAttributes:
     MILVUS_UPSERT_COLLECTION_NAME = "db.milvus.upsert.collection_name"
     MILVUS_UPSERT_DATA_COUNT = "db.milvus.upsert.data_count"
     MILVUS_UPSERT_PARTITION_NAME = "db.milvus.upsert.partition_name"
-    MILVUS_UPSERT_TIMEOUT_COUNT = "db.milvus.upsert.timeout_count"
+    MILVUS_UPSERT_TIMEOUT = "db.milvus.upsert.timeout"
 
     # Qdrant
     QDRANT_SEARCH_COLLECTION_NAME = "qdrant.search.collection_name"