enrich search

Author: alizenhom

src/examples/milvus_example/main.py

@@ -56,6 +56,10 @@ def vector_search(collection_name: str = COLLECTION_NAME, queries: List[str] = [
         data=query_vectors,  # query vectors
         limit=2,  # number of returned entities
         output_fields=["text", "subject"],  # specifies fields to be returned
+        timeout=10,
+        partition_names=["history"],
+        anns_field="vector",
+        search_params={"nprobe": 10},
     )
 
 

src/langtrace_python_sdk/instrumentation/milvus/patch.py

@@ -1,10 +1,12 @@
+from langtrace_python_sdk.utils.silently_fail import silently_fail
 from opentelemetry.trace import Tracer
 from opentelemetry.trace import SpanKind
-from langtrace_python_sdk.utils import handle_span_error
+from langtrace_python_sdk.utils import handle_span_error, set_span_attribute
 from langtrace_python_sdk.utils.llm import (
     get_extra_attributes,
     set_span_attributes,
 )
+import json
 
 
 def generic_patch(api, version: str, tracer: Tracer):
@@ -38,6 +40,8 @@ def traced_method(wrapped, instance, args, kwargs):
                 if operation == "query":
                     set_query_response_attributes(span, result)
 
+                if operation == "search":
+                    set_search_response_attributes(span, result)
                 return result
             except Exception as err:
                 handle_span_error(span, err)
@@ -46,10 +50,12 @@ def traced_method(wrapped, instance, args, kwargs):
     return traced_method
 
 
+@silently_fail
 def set_create_collection_attributes(span_attributes, kwargs):
     span_attributes["db.dimension"] = kwargs.get("dimension", None)
 
 
+@silently_fail
 def set_insert_or_upsert_attributes(span_attributes, kwargs):
     data = kwargs.get("data")
     timeout = kwargs.get("timeout")
@@ -60,6 +66,7 @@ def set_insert_or_upsert_attributes(span_attributes, kwargs):
     span_attributes["db.partition_name"] = partition_name
 
 
+@silently_fail
 def set_search_attributes(span_attributes, kwargs):
     data = kwargs.get("data")
     filter = kwargs.get("filter")
@@ -69,17 +76,17 @@ def set_search_attributes(span_attributes, kwargs):
     timeout = kwargs.get("timeout")
     partition_names = kwargs.get("partition_names")
     anns_field = kwargs.get("anns_field")
-
     span_attributes["db.num_queries"] = len(data) if data else None
     span_attributes["db.filter"] = filter
     span_attributes["db.limit"] = limit
-    span_attributes["db.output_fields"] = output_fields
-    span_attributes["db.search_params"] = search_params
-    span_attributes["db.partition_names"] = partition_names
+    span_attributes["db.output_fields"] = json.dumps(output_fields)
+    span_attributes["db.search_params"] = json.dumps(search_params)
+    span_attributes["db.partition_names"] = json.dumps(partition_names)
     span_attributes["db.anns_field"] = anns_field
     span_attributes["db.timeout"] = timeout
 
 
+@silently_fail
 def set_query_attributes(span_attributes, kwargs):
     filter = kwargs.get("filter")
     output_fields = kwargs.get("output_fields")
@@ -94,9 +101,25 @@ def set_query_attributes(span_attributes, kwargs):
     span_attributes["db.ids"] = ids
 
 
+@silently_fail
 def set_query_response_attributes(span, result):
+    set_span_attribute(span, name="db.num_matches", value=len(result))
     for match in result:
         span.add_event(
             "db.query.match",
             attributes=match,
         )
+
+
+@silently_fail
+def set_search_response_attributes(span, result):
+    for res in result:
+        for match in res:
+            span.add_event(
+                "db.search.match",
+                attributes={
+                    "id": match["id"],
+                    "distance": str(match["distance"]),
+                    "entity": json.dumps(match["entity"]),
+                },
+            )