feat(milvus): Add error.type attribute from OpenTelemetry Semantic Conventions (#3009)

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>
Co-authored-by: Nir Gazit <[email protected]>

Author: 3 people

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

@@ -1,13 +1,24 @@
 from opentelemetry.instrumentation.milvus.utils import dont_throw
 from opentelemetry.semconv.trace import SpanAttributes
-
+from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
 from opentelemetry import context as context_api
 from opentelemetry.instrumentation.utils import (
     _SUPPRESS_INSTRUMENTATION_KEY,
 )
 from opentelemetry.semconv_ai import Events, EventAttributes
 from opentelemetry.semconv_ai import SpanAttributes as AISpanAttributes
 
+from pymilvus.client.types import Status
+from pymilvus.exceptions import ErrorCode
+
+code_to_error_type = {}
+for name in dir(Status):
+    value = getattr(Status, name)
+    if isinstance(value, int):
+        code_to_error_type[value] = name
+
+code_to_error_type.update({code.value: code.name for code in ErrorCode})
+
 
 def _with_tracer_wrapper(func):
     """Helper for providing tracer for wrapper functions."""
@@ -56,16 +67,22 @@ def _wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
         elif to_wrap.get("method") == "hybrid_search":
             _set_hybrid_search_attributes(span, kwargs)
 
-        return_value = wrapped(*args, **kwargs)
-
-        if to_wrap.get("method") == "query":
-            _add_query_result_events(span, return_value)
-
-        if (
-            to_wrap.get("method") == "search"
-            or to_wrap.get("method") == "hybrid_search"
-        ):
-            _add_search_result_events(span, return_value)
+        try:
+            return_value = wrapped(*args, **kwargs)
+            if to_wrap.get("method") == "query":
+                _add_query_result_events(span, return_value)
+
+            if (
+                to_wrap.get("method") == "search"
+                or to_wrap.get("method") == "hybrid_search"
+            ):
+                _add_search_result_events(span, return_value)
+        except Exception as e:
+            error_type = code_to_error_type.get(
+                getattr(e, "code", None), type(e).__name__
+            )
+            span.set_attribute(ERROR_TYPE, error_type)
+            raise
 
     return return_value
 
@@ -345,6 +362,7 @@ def set_global_stats():
         _set_span_attribute(
             span, AISpanAttributes.MILVUS_SEARCH_RESULT_COUNT, total_matches
         )
+
     for query_idx, query_results in enumerate(kwargs):
 
         query_distances = []

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

@@ -0,0 +1,91 @@
+import os
+import random
+
+import pymilvus
+import pytest
+from opentelemetry.trace.status import StatusCode
+
+path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "milvus.db")
+milvus = pymilvus.MilvusClient(uri=path)
+
+
+@pytest.fixture
+def collection():
+    collection_name = "Colors"
+    milvus.create_collection(collection_name=collection_name, dimension=5)
+    yield collection_name
+    milvus.drop_collection(collection_name=collection_name)
+
+
+def insert_data(collection):
+    colors = [
+        "green",
+        "blue",
+        "yellow",
+        "red",
+        "black",
+        "white",
+        "purple",
+        "pink",
+        "orange",
+        "grey",
+    ]
+    data = [
+        {
+            "id": i,
+            "vector": [random.uniform(-1, 1) for _ in range(5)],
+            "color": random.choice(colors),
+            "tag": random.randint(1000, 9999),
+        }
+        for i in range(1000)
+    ]
+    data += [
+        {
+            "id": 1000,
+            "vector": [random.uniform(-1, 1) for _ in range(5)],
+            "color": "brown",
+            "tag": 1234,
+        },
+        {
+            "id": 1001,
+            "vector": [random.uniform(-1, 1) for _ in range(5)],
+            "color": "brown",
+            "tag": 5678,
+        },
+        {
+            "id": 1002,
+            "vector": [random.uniform(-1, 1) for _ in range(5)],
+            "color": "brown",
+            "tag": 9101,
+        },
+    ]
+    for i in data:
+        i["color_tag"] = "{}_{}".format(i["color"], i["tag"])
+    milvus.insert(collection_name=collection, data=data)
+
+
+def test_milvus_single_vector_search(exporter, collection):
+    insert_data(collection)
+
+    query_vector = [random.uniform(-1, 1) for _ in range(5)]
+    search_params = {"radius": 0.5, "metric_type": "COSINE", "index_type": "IVF_FLAT"}
+    with pytest.raises(Exception):
+        milvus.search(
+            collection_name="random",  # non-existent collection
+            data=[query_vector],
+            anns_field="vector",
+            search_params=search_params,
+            output_fields=["color_tag"],
+            limit=3,
+            timeout=10,
+        )
+
+    # Get finished spans
+    spans = exporter.get_finished_spans()
+    span = next(span for span in spans if span.name == "milvus.search")
+
+    # Check if status code is error
+    assert span.status.status_code == StatusCode.ERROR
+
+    # Check the span attributes related to search
+    assert span.attributes.get("error.type") == "COLLECTION_NOT_FOUND"