mongodb kickstart + openai embedding enrichment

Author: alizenhom

src/examples/mongo_vector_search_example/main.py

@@ -0,0 +1,63 @@
+import pymongo
+import os
+from dotenv import load_dotenv
+from openai import OpenAI
+from langtrace_python_sdk import langtrace
+
+load_dotenv()
+langtrace.init()
+MODEL = "text-embedding-ada-002"
+openai_client = OpenAI()
+client = pymongo.MongoClient(os.environ["MONGO_URI"])
+
+
+# Define a function to generate embeddings
+def get_embedding(text):
+    """Generates vector embeddings for the given text."""
+    embedding = (
+        openai_client.embeddings.create(input=[text], model=MODEL).data[0].embedding
+    )
+    return embedding
+
+
+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,
+                "limit": 10,
+            }
+        },
+        {
+            "$project": {
+                "_id": 0,
+                "plot": 1,
+                "title": 1,
+                "score": {"$meta": "vectorSearchScore"},
+            }
+        },
+    ]
+
+    result = embedded_movies_collection.aggregate(pipeline)
+    for doc in result:
+        print(doc)
+
+
+if __name__ == "__main__":
+    try:
+        vector_query()
+    except Exception as e:
+        print(e)
+    finally:
+        client.close()

src/langtrace_python_sdk/instrumentation/__init__.py

@@ -20,6 +20,7 @@
 from .mistral import MistralInstrumentation
 from .embedchain import EmbedchainInstrumentation
 from .litellm import LiteLLMInstrumentation
+from .pymongo import PyMongoInstrumentation
 
 __all__ = [
     "AnthropicInstrumentation",
@@ -44,4 +45,5 @@
     "VertexAIInstrumentation",
     "GeminiInstrumentation",
     "MistralInstrumentation",
+    "PyMongoInstrumentation",
 ]

src/langtrace_python_sdk/instrumentation/openai/patch.py

@@ -27,6 +27,7 @@
     set_event_completion,
     StreamWrapper,
     set_span_attributes,
+    set_usage_attributes,
 )
 from langtrace_python_sdk.types import NOT_GIVEN
 
@@ -446,6 +447,14 @@ def traced_method(
             span_attributes[SpanAttributes.LLM_REQUEST_EMBEDDING_INPUTS] = json.dumps(
                 [kwargs.get("input", "")]
             )
+            span_attributes[SpanAttributes.LLM_PROMPTS] = json.dumps(
+                [
+                    {
+                        "role": "user",
+                        "content": kwargs.get("input"),
+                    }
+                ]
+            )
 
         attributes = LLMSpanAttributes(**filter_valid_attributes(span_attributes))
 
@@ -459,6 +468,11 @@ def traced_method(
             try:
                 # Attempt to call the original method
                 result = wrapped(*args, **kwargs)
+                usage = getattr(result, "usage", None)
+                if usage:
+                    set_usage_attributes(
+                        span, {"prompt_tokens": getattr(usage, "prompt_tokens", 0)}
+                    )
                 span.set_status(StatusCode.OK)
                 return result
             except Exception as err:

src/langtrace_python_sdk/instrumentation/pymongo/__init__.py

@@ -0,0 +1,5 @@
+from .instrumentation import PyMongoInstrumentation
+
+__all__ = [
+    "PyMongoInstrumentation",
+]

src/langtrace_python_sdk/instrumentation/pymongo/instrumentation.py

@@ -0,0 +1,45 @@
+"""
+Copyright (c) 2024 Scale3 Labs
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+
+from typing import Collection
+from importlib_metadata import version as v
+from wrapt import wrap_function_wrapper as _W
+from .patch import generic_patch
+
+
+class PyMongoInstrumentation(BaseInstrumentor):
+    """
+    The PyMongoInstrumentation class represents the PyMongo instrumentation
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["pymongo >= 4.0.0"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = v("pymongo")
+        _W(
+            module="pymongo.collection",
+            name="Collection.find",
+            wrapper=generic_patch(version, tracer),
+        )
+
+    def _uninstrument(self, **kwargs):
+        pass

src/langtrace_python_sdk/instrumentation/pymongo/patch.py

@@ -0,0 +1,8 @@
+def generic_patch(version, tracer):
+    def traced_method(wrapped, instance, args, kwargs):
+        print("kwargs", kwargs)
+        print("args", args)
+        print("instance", instance)
+        wrapped(*args, **kwargs)
+
+    return traced_method

src/langtrace_python_sdk/langtrace.py

@@ -58,6 +58,7 @@
     AutogenInstrumentation,
     VertexAIInstrumentation,
     WeaviateInstrumentation,
+    PyMongoInstrumentation,
 )
 from langtrace_python_sdk.types import DisableInstrumentations, InstrumentationMethods
 from langtrace_python_sdk.utils import (
@@ -230,6 +231,7 @@ def init(
         "google-generativeai": GeminiInstrumentation(),
         "mistralai": MistralInstrumentation(),
         "autogen": AutogenInstrumentation(),
+        "pymongo": PyMongoInstrumentation(),
     }
 
     init_instrumentations(config.disable_instrumentations, all_instrumentations)