langtrace-python-sdk 3.3.1__py3-none-any.whl → 3.3.3__py3-none-any.whl

examples/mongo_vector_search_example/main.py

@@ -0,0 +1,61 @@
+from langtrace_python_sdk import langtrace, with_langtrace_root_span
+import pymongo
+import os
+from dotenv import load_dotenv
+from openai import OpenAI
+
+load_dotenv()
+langtrace.init(write_spans_to_console=False, batch=False)
+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
+
+
+@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("time travel"),
+                "numCandidates": 150,
+                "limit": 10,
+            }
+        },
+        {
+            "$project": {
+                "_id": 0,
+                "plot": 1,
+                "title": 1,
+                "score": {"$meta": "vectorSearchScore"},
+            }
+        },
+    ]
+
+    result = embedded_movies_collection.aggregate(pipeline)
+    for doc in result:
+        # print(doc)
+        pass
+
+
+if __name__ == "__main__":
+    try:
+        vector_query()
+    except Exception as e:
+        print("error", e)
+    finally:
+        client.close()

langtrace_python_sdk/constants/instrumentation/common.py

@@ -34,6 +34,7 @@ SERVICE_PROVIDERS = {
     "EMBEDCHAIN": "Embedchain",
     "AUTOGEN": "Autogen",
     "XAI": "XAI",
+    "MONGODB": "MongoDB",
     "AWS_BEDROCK": "AWS Bedrock",
     "CEREBRAS": "Cerebras",
 }

langtrace_python_sdk/constants/instrumentation/pymongo.py

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

langtrace_python_sdk/instrumentation/__init__.py

@@ -21,6 +21,7 @@ from .mistral import MistralInstrumentation
 from .aws_bedrock import AWSBedrockInstrumentation
 from .embedchain import EmbedchainInstrumentation
 from .litellm import LiteLLMInstrumentation
+from .pymongo import PyMongoInstrumentation
 from .cerebras import CerebrasInstrumentation
 
 __all__ = [
@@ -46,6 +47,7 @@ __all__ = [
     "VertexAIInstrumentation",
     "GeminiInstrumentation",
     "MistralInstrumentation",
+    "PyMongoInstrumentation",
     "AWSBedrockInstrumentation",
     "CerebrasInstrumentation",
 ]

langtrace_python_sdk/instrumentation/openai/patch.py

@@ -27,6 +27,7 @@ from langtrace_python_sdk.utils.llm import (
     set_event_completion,
     StreamWrapper,
     set_span_attributes,
+    set_usage_attributes,
 )
 from langtrace_python_sdk.types import NOT_GIVEN
 
@@ -450,6 +451,14 @@ def embeddings_create(version: str, tracer: Tracer) -> Callable:
             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))
 
@@ -463,6 +472,11 @@ def embeddings_create(version: str, tracer: Tracer) -> Callable:
             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:

langtrace_python_sdk/instrumentation/pymongo/__init__.py

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

langtrace_python_sdk/instrumentation/pymongo/instrumentation.py

@@ -0,0 +1,47 @@
+"""
+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
+from langtrace_python_sdk.constants.instrumentation.pymongo import APIS
+
+
+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")
+        for api in APIS.values():
+            _W(
+                module=api["MODULE"],
+                name=api["METHOD"],
+                wrapper=generic_patch(api["SPAN_NAME"], version, tracer),
+            )
+
+    def _uninstrument(self, **kwargs):
+        pass

langtrace_python_sdk/instrumentation/pymongo/patch.py

@@ -0,0 +1,66 @@
+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, handle_span_error
+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):
+        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:
+                handle_span_error(span, err)
+                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))

langtrace_python_sdk/langtrace.py

@@ -64,6 +64,7 @@ from langtrace_python_sdk.instrumentation import (
     AutogenInstrumentation,
     VertexAIInstrumentation,
     WeaviateInstrumentation,
+    PyMongoInstrumentation,
     CerebrasInstrumentation,
 )
 from opentelemetry.util.re import parse_env_headers
@@ -79,8 +80,6 @@ from langtrace_python_sdk.utils.langtrace_sampler import LangtraceSampler
 from langtrace_python_sdk.extensions.langtrace_exporter import LangTraceExporter
 from sentry_sdk.types import Event, Hint
 
-logging.disable(level=logging.INFO)
-
 
 class LangtraceConfig:
     def __init__(self, **kwargs):
@@ -216,7 +215,7 @@ def init(
     disable_logging: bool = False,
     headers: Dict[str, str] = {},
 ):
-    logging.disable(level=logging.INFO)
+
     check_if_sdk_is_outdated()
     config = LangtraceConfig(
         api_key=api_key,
@@ -232,6 +231,7 @@ def init(
     )
 
     if config.disable_logging:
+        logging.disable(level=logging.INFO)
         sys.stdout = open(os.devnull, "w")
 
     host = get_host(config)
@@ -282,6 +282,7 @@ def init(
         "mistralai": MistralInstrumentation(),
         "boto3": AWSBedrockInstrumentation(),
         "autogen": AutogenInstrumentation(),
+        "pymongo": PyMongoInstrumentation(),
         "cerebras-cloud-sdk": CerebrasInstrumentation(),
     }
 

langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "3.3.1"
+__version__ = "3.3.3"

{langtrace_python_sdk-3.3.1.dist-info → langtrace_python_sdk-3.3.3.dist-info}/METADATA

@@ -1,11 +1,9 @@
 Metadata-Version: 2.3
 Name: langtrace-python-sdk
-Version: 3.3.1
+Version: 3.3.3
 Summary: Python SDK for LangTrace
 Project-URL: Homepage, https://github.com/Scale3-Labs/langtrace-python-sdk
 Author-email: Scale3 Labs <engineering@scale3labs.com>
-License-Expression: Apache-2.0
-License-File: LICENSE
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3

{langtrace_python_sdk-3.3.1.dist-info → langtrace_python_sdk-3.3.3.dist-info}/RECORD

@@ -72,6 +72,7 @@ examples/mistral_example/__init__.py,sha256=4kMzYY7whzUFkT9LvqPp9wSdwNtk8pq28pmZ
 examples/mistral_example/complete.py,sha256=Ydf5iOCGM2tySSd5vzGvNh3Qz39aTpxcCItBii9b0ig,589
 examples/mistral_example/complete_async.py,sha256=-SVOLNLNi5JL3o4obBsCXKOt8A6b61j3otbrrYPOOtM,586
 examples/mistral_example/embeddings.py,sha256=LKq_k3Y-TS2SCkyFKw2tLiYreVRHqlY6z0Gfh1y_7u0,468
+examples/mongo_vector_search_example/main.py,sha256=6bHOxR3LfAZwmGmv74jgNCciATS5h5e-Tb5z0cD0z3Q,1542
 examples/ollama_example/__init__.py,sha256=qOx0jGCPuSpRCPiqtDVm7F0z8hIZ8C75hDZ_C8Apz-s,399
 examples/ollama_example/basic.py,sha256=EPbsigOF4xBDBgLgAD0EzPo737ycVm7aXZr7F5Xt-A4,1062
 examples/openai_example/__init__.py,sha256=6faH7wTegSozKmS89sd1Tgv8AcEH0GfKkC7YaBWA8tg,849
@@ -102,8 +103,8 @@ examples/vertexai_example/main.py,sha256=gndId5X5ksD-ycxnAWMdEqIDbLc3kz5Vt8vm4YP
 examples/weaviate_example/__init__.py,sha256=8JMDBsRSEV10HfTd-YC7xb4txBjD3la56snk-Bbg2Kw,618
 examples/weaviate_example/query_text.py,sha256=wPHQTc_58kPoKTZMygVjTj-2ZcdrIuaausJfMxNQnQc,127162
 langtrace_python_sdk/__init__.py,sha256=VZM6i71NR7pBQK6XvJWRelknuTYUhqwqE7PlicKa5Wg,1166
-langtrace_python_sdk/langtrace.py,sha256=IBl7tapb5Mig5362WdIoZUmYWE-ZuqFe2Bk4IJxk-Xs,12488
-langtrace_python_sdk/version.py,sha256=XErusAMGUPwBUpdA6BLyq8CjU-6n6gLlBRymgSC8Y-0,22
+langtrace_python_sdk/langtrace.py,sha256=rxd-Vn0hTqqd02E5S6KzwjiTdSMttNfhY63d74Xi_gk,12529
+langtrace_python_sdk/version.py,sha256=Au2yR9CeJDYeCsUV7lIDllPVwz8EynN-QwVbxO3RWOM,22
 langtrace_python_sdk/constants/__init__.py,sha256=3CNYkWMdd1DrkGqzLUgNZXjdAlM6UFMlf_F-odAToyc,146
 langtrace_python_sdk/constants/exporter/langtrace_exporter.py,sha256=d-3Qn5C_NTy1NkmdavZvy-6vePwTC5curN6QMy2haHc,50
 langtrace_python_sdk/constants/instrumentation/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -111,7 +112,7 @@ langtrace_python_sdk/constants/instrumentation/anthropic.py,sha256=YX3llt3zwDY6X
 langtrace_python_sdk/constants/instrumentation/aws_bedrock.py,sha256=f9eukqoxrPgPeaBJX2gpBUz1uu0dZIPahOpvoudfbH8,310
 langtrace_python_sdk/constants/instrumentation/chroma.py,sha256=hiPGYdHS0Yj4Kh3eaYBbuCAl_swqIygu80yFqkOgdak,955
 langtrace_python_sdk/constants/instrumentation/cohere.py,sha256=tf9sDfb5K3qOAHChEE5o8eYWPZ1io58VsOjZDCZPxfw,577
-langtrace_python_sdk/constants/instrumentation/common.py,sha256=gs7xupRbUW_4UjXFnmqiLGRktA8lllXHRcOLeQFKek4,1125
+langtrace_python_sdk/constants/instrumentation/common.py,sha256=OJcZBJptWlZGmsdI47n9pqH4q-KkYB9xN0WvtQfeaoo,1151
 langtrace_python_sdk/constants/instrumentation/embedchain.py,sha256=HodCJvaFjILoOG50OwFObxfVxt_8VUaIAIqvgoN3tzo,278
 langtrace_python_sdk/constants/instrumentation/gemini.py,sha256=UAmfgg9FM7uNeOCdPfWlir6OIH-8BoxFGPRpdBd9ZZs,358
 langtrace_python_sdk/constants/instrumentation/groq.py,sha256=VFXmIl4aqGY_fS0PAmjPj_Qm7Tibxbx7Ur_e7rQpqXc,134
@@ -120,13 +121,14 @@ langtrace_python_sdk/constants/instrumentation/mistral.py,sha256=9PlmcC5P5_BHJ-z
 langtrace_python_sdk/constants/instrumentation/ollama.py,sha256=H_-S0xjqRsi5qSp7mAlK7Y9NlQ3BqOkG6ASogqqgdJY,212
 langtrace_python_sdk/constants/instrumentation/openai.py,sha256=uEOH5UXapU2DSf2AdgXTRhhJEHGWXUNFkUGD5QafflM,1164
 langtrace_python_sdk/constants/instrumentation/pinecone.py,sha256=0TityERbGWaHGSN8-vyYZtYCjVj8fQOKae8lng0O0Bk,478
+langtrace_python_sdk/constants/instrumentation/pymongo.py,sha256=bZR3anvH2OjQCj6ZCLMmdJcWcmlHhEGezS_3sx9RyJ8,195
 langtrace_python_sdk/constants/instrumentation/qdrant.py,sha256=yL7BopNQTXW7L7Z-gVM2PdusKD7r9qqcATvczFd7NtQ,1999
 langtrace_python_sdk/constants/instrumentation/vertexai.py,sha256=0s2vX3Y0iwjOPkUg5lAKi-7o3LaNivDSBBbF-o695Ok,1266
 langtrace_python_sdk/constants/instrumentation/weaviate.py,sha256=gtv-JBxvNGClEMxClmRKzjJ1khgOonsli4D_k9IagSE,2601
 langtrace_python_sdk/extensions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langtrace_python_sdk/extensions/langtrace_exporter.py,sha256=UFupNL03zklVd5penpsfXjbWSb5qB39mEv2BY2wczSs,6307
 langtrace_python_sdk/extensions/langtrace_filesystem.py,sha256=34fZutG28EJ66l67OvTGsydAH3ZpXgikdE7hVLqBpG4,7863
-langtrace_python_sdk/instrumentation/__init__.py,sha256=0wSimp9RM2Z8OpGW-h-ftDCWjXDcer0UAWg-BYXyZ0s,1857
+langtrace_python_sdk/instrumentation/__init__.py,sha256=sH9bZVQN9oHZHjGnIFLSue3jbkhvHTfb4D6x3nuChCQ,1931
 langtrace_python_sdk/instrumentation/anthropic/__init__.py,sha256=donrurJAGYlxrSRA3BIf76jGeUcAx9Tq8CVpah68S0Y,101
 langtrace_python_sdk/instrumentation/anthropic/instrumentation.py,sha256=ndXdruI0BG7n75rsuEpKjfzePxrZxg40gZ39ONmD_v4,1845
 langtrace_python_sdk/instrumentation/anthropic/patch.py,sha256=ztPN4VZujoxYOKhTbFnup7Ibms9NAzYCPAJY43NUgKw,4935
@@ -188,11 +190,14 @@ langtrace_python_sdk/instrumentation/ollama/instrumentation.py,sha256=jdsvkqUJAA
 langtrace_python_sdk/instrumentation/ollama/patch.py,sha256=7ETx0tQic5h_kH1f-IeptFwgNTBb4hSkTkWsB18Avm0,5375
 langtrace_python_sdk/instrumentation/openai/__init__.py,sha256=VPHRNCQEdkizIVP2d0Uw_a7t8XOTSTprEIB8oboJFbs,95
 langtrace_python_sdk/instrumentation/openai/instrumentation.py,sha256=PZxI0qfoud1VsKGmJu49YDp0Z9z9TzCR8qxR3uznOMA,2810
-langtrace_python_sdk/instrumentation/openai/patch.py,sha256=MacQoaqs5b3zgCkD6z2k6fP9mruRtLSLWNNBXuZ70m8,24431
+langtrace_python_sdk/instrumentation/openai/patch.py,sha256=GINHfq5HHSfOzMr1HErwLuzOFx5pzmclE9Pwa4ziKBQ,24946
 langtrace_python_sdk/instrumentation/openai/types.py,sha256=aVkoa7tmAbYfyOhnyMrDaVjQuwhmRNLMthlNtKMtWX8,4311
 langtrace_python_sdk/instrumentation/pinecone/__init__.py,sha256=DzXyGh9_MGWveJvXULkFwdkf7PbG2s3bAWtT1Dmz7Ok,99
 langtrace_python_sdk/instrumentation/pinecone/instrumentation.py,sha256=HDXkRITrVPwdQEoOYJOfMzZE_2-vDDvuqHTlD8W1lQw,1845
 langtrace_python_sdk/instrumentation/pinecone/patch.py,sha256=MIbUcEsVzl4W_pfq81LP9QFVhwPB8rHF0Aod9pq_j-o,5249
+langtrace_python_sdk/instrumentation/pymongo/__init__.py,sha256=Q_ecuvZ0j5TxOCdD69BfZhCPT6rj7w5lHbTuDSpLYD4,97
+langtrace_python_sdk/instrumentation/pymongo/instrumentation.py,sha256=fG938JR2Be3zy19JUyvKxafYQPavnUV9eMtD7Vn0_fc,1597
+langtrace_python_sdk/instrumentation/pymongo/patch.py,sha256=4K-_QWTXoWDzpYEgaWPkf6kvXGnAmHKvp8H2gMmPokA,2330
 langtrace_python_sdk/instrumentation/qdrant/__init__.py,sha256=TaIGSAEPysrL23KJ5FcEL1tfQogrKCtEQ75_u62eqso,95
 langtrace_python_sdk/instrumentation/qdrant/instrumentation.py,sha256=vl2eKSP55aqDo1JiRlvOUBrr6kddvG9Z5dCYew2OG08,1816
 langtrace_python_sdk/instrumentation/qdrant/patch.py,sha256=IgdozFyKqB8n72BjKvBDiMhYM4o75DReD0I8_uIQ7KY,5015
@@ -254,8 +259,8 @@ tests/pinecone/cassettes/test_query.yaml,sha256=b5v9G3ssUy00oG63PlFUR3JErF2Js-5A
 tests/pinecone/cassettes/test_upsert.yaml,sha256=neWmQ1v3d03V8WoLl8FoFeeCYImb8pxlJBWnFd_lITU,38607
 tests/qdrant/conftest.py,sha256=9n0uHxxIjWk9fbYc4bx-uP8lSAgLBVx-cV9UjnsyCHM,381
 tests/qdrant/test_qdrant.py,sha256=pzjAjVY2kmsmGfrI2Gs2xrolfuaNHz7l1fqGQCjp5_o,3353
-langtrace_python_sdk-3.3.1.dist-info/METADATA,sha256=S13-jNowYfopmvPCqaqKnazad-r6iOp_DNop-VrgAps,15996
-langtrace_python_sdk-3.3.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-langtrace_python_sdk-3.3.1.dist-info/entry_points.txt,sha256=1_b9-qvf2fE7uQNZcbUei9vLpFZBbbh9LrtGw95ssAo,70
-langtrace_python_sdk-3.3.1.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-langtrace_python_sdk-3.3.1.dist-info/RECORD,,
+langtrace_python_sdk-3.3.3.dist-info/METADATA,sha256=dWEl_fFJz9n8MjBCDuafoksbyiaCxKRnphD8La5xtyM,15943
+langtrace_python_sdk-3.3.3.dist-info/WHEEL,sha256=3U_NnUcV_1B1kPkYaPzN-irRckL5VW_lytn0ytO_kRY,87
+langtrace_python_sdk-3.3.3.dist-info/entry_points.txt,sha256=1_b9-qvf2fE7uQNZcbUei9vLpFZBbbh9LrtGw95ssAo,70
+langtrace_python_sdk-3.3.3.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+langtrace_python_sdk-3.3.3.dist-info/RECORD,,

{langtrace_python_sdk-3.3.1.dist-info → langtrace_python_sdk-3.3.3.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.25.0
+Generator: hatchling 1.26.1
 Root-Is-Purelib: true
 Tag: py3-none-any