mongodb vector search instrument

Author: alizenhom

src/examples/mongo_vector_search_example/main.py

@@ -1,11 +1,11 @@
+from langtrace_python_sdk import langtrace, with_langtrace_root_span
 import pymongo
 import os
 from dotenv import load_dotenv
 from openai import OpenAI
-from langtrace_python_sdk import langtrace
 
 load_dotenv()
-langtrace.init()
+langtrace.init(write_spans_to_console=False, batch=False)
 MODEL = "text-embedding-ada-002"
 openai_client = OpenAI()
 client = pymongo.MongoClient(os.environ["MONGO_URI"])
@@ -20,22 +20,19 @@ def get_embedding(text):
     return embedding
 
 
+@with_langtrace_root_span("mongo-vector-search")
 def vector_query():
     db = client["sample_mflix"]
 
     embedded_movies_collection = db["embedded_movies"]
-
     # define pipeline
     pipeline = [
         {
             "$vectorSearch": {
                 "index": "vector_index",
                 "path": "plot_embedding",
-                "queryVector": get_embedding(
-                    "A movie about a hacker that had a really rough childhood and been trying to convince his father otherwise."
-                ),
-                # "numCandidates": 150,
-                "exact": True,
+                "queryVector": get_embedding("time travel"),
+                "numCandidates": 150,
                 "limit": 10,
             }
         },
@@ -51,13 +48,14 @@ def vector_query():
 
     result = embedded_movies_collection.aggregate(pipeline)
     for doc in result:
-        print(doc)
+        # print(doc)
+        pass
 
 
 if __name__ == "__main__":
     try:
         vector_query()
     except Exception as e:
-        print(e)
+        print("error", e)
     finally:
         client.close()

src/langtrace_python_sdk/constants/instrumentation/common.py

@@ -34,6 +34,7 @@
     "EMBEDCHAIN": "Embedchain",
     "AUTOGEN": "Autogen",
     "XAI": "XAI",
+    "MONGODB": "MongoDB",
 }
 
 LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY = "langtrace_additional_attributes"

src/langtrace_python_sdk/constants/instrumentation/pymongo.py

@@ -1,6 +1,8 @@
 APIS = {
     "AGGREGATE": {
-        "METHOD": "aggregate",
+        "MODULE": "pymongo.collection",
+        "METHOD": "Collection.aggregate",
         "OPERATION": "aggregate",
+        "SPAN_NAME": "MongoDB Aggregate",
     },
 }

src/langtrace_python_sdk/instrumentation/pymongo/instrumentation.py

@@ -38,9 +38,9 @@ def _instrument(self, **kwargs):
         version = v("pymongo")
         for api in APIS.values():
             _W(
-                module="pymongo.collection",
-                name=f"Collection.{api['METHOD']}",
-                wrapper=generic_patch(version, tracer),
+                module=api["MODULE"],
+                name=api["METHOD"],
+                wrapper=generic_patch(api["SPAN_NAME"], version, tracer),
             )
 
     def _uninstrument(self, **kwargs):

src/langtrace_python_sdk/instrumentation/pymongo/patch.py

@@ -1,5 +1,73 @@
-def generic_patch(version, tracer):
+from langtrace_python_sdk.utils.llm import (
+    get_langtrace_attributes,
+    get_span_name,
+    set_span_attributes,
+    set_span_attribute,
+)
+from langtrace_python_sdk.utils import deduce_args_and_kwargs
+from opentelemetry.trace.status import Status, StatusCode
+from opentelemetry.trace import SpanKind
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
+from langtrace.trace_attributes import DatabaseSpanAttributes
+
+import json
+
+
+def generic_patch(name, version, tracer):
     def traced_method(wrapped, instance, args, kwargs):
-        wrapped(*args, **kwargs)
+        database = instance.database.__dict__
+        span_attributes = {
+            **get_langtrace_attributes(
+                version=version,
+                service_provider=SERVICE_PROVIDERS["MONGODB"],
+                vendor_type="vectordb",
+            ),
+            "db.system": "mongodb",
+            "db.query": "aggregate",
+        }
+
+        attributes = DatabaseSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(
+            get_span_name(name), kind=SpanKind.CLIENT
+        ) as span:
+            if span.is_recording():
+                set_input_attributes(
+                    span, deduce_args_and_kwargs(wrapped, *args, **kwargs)
+                )
+                set_span_attributes(span, attributes)
+
+            try:
+                result = wrapped(*args, **kwargs)
+                print(result)
+                for doc in result:
+                    if span.is_recording():
+                        span.add_event(
+                            name="db.query.match",
+                            attributes={**doc},
+                        )
+                return result
+            except Exception as err:
+                # Record the exception in the span
+                span.record_exception(err)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(err)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
 
     return traced_method
+
+
+def set_input_attributes(span, args):
+    pipeline = args.get("pipeline", None)
+    for stage in pipeline:
+        for k, v in stage.items():
+            if k == "$vectorSearch":
+                set_span_attribute(span, "db.index", v.get("index", None))
+                set_span_attribute(span, "db.path", v.get("path", None))
+                set_span_attribute(span, "db.top_k", v.get("numCandidates"))
+                set_span_attribute(span, "db.limit", v.get("limit"))
+            else:
+                set_span_attribute(span, k, json.dumps(v))