langtrace-python-sdk 2.0.6__py3-none-any.whl → 2.0.8__py3-none-any.whl

examples/pinecone_example/basic.py

@@ -2,19 +2,16 @@
 This example demonstrates how to use Pinecone with Langtrace.
 """
 
+from langtrace_python_sdk import langtrace
 from dotenv import find_dotenv, load_dotenv
 from openai import OpenAI
 from pinecone import Pinecone
 
-from langtrace_python_sdk import langtrace
 from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span
 
 _ = load_dotenv(find_dotenv())
 
-langtrace.init(
-    write_spans_to_console=True,
-    disable_instrumentations={"all_except": ["pinecone", "openai"]},
-)
+langtrace.init()
 
 client = OpenAI()
 pinecone = Pinecone()
@@ -30,10 +27,14 @@ def basic():
 
     embedding = result.data[0].embedding
 
-    unique_id = "randomid"
+    unique_id = "unique_random_id"
     data_to_upsert = {"id": unique_id, "values": embedding}
 
     index = pinecone.Index("test-index")
-    index.upsert(vectors=[data_to_upsert])
+    res = index.upsert(vectors=[data_to_upsert], namespace="test-namespace")
+    print("res", res)
 
-    resp = index.query(vector=embedding, top_k=1)
+    resp = index.query(
+        vector=embedding, top_k=1, include_values=False, namespace="test-namespace"
+    )
+    print(resp)

langtrace_python_sdk/__init__.py

@@ -16,5 +16,6 @@ limitations under the License.
 
 from langtrace_python_sdk import langtrace
 from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span
+from langtrace_python_sdk.utils.prompt_registry import get_prompt_from_registry
 
-__all__ = ["langtrace", "with_langtrace_root_span"]
+__all__ = ["langtrace", "with_langtrace_root_span", "get_prompt_from_registry"]

langtrace_python_sdk/instrumentation/pinecone/patch.py

@@ -15,6 +15,8 @@ limitations under the License.
 """
 
 from langtrace.trace_attributes import DatabaseSpanAttributes
+from langtrace_python_sdk.utils import set_span_attribute
+from langtrace_python_sdk.utils.silently_fail import silently_fail
 from opentelemetry import baggage
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
@@ -24,6 +26,7 @@ from langtrace_python_sdk.constants.instrumentation.common import (
     SERVICE_PROVIDERS,
 )
 from langtrace_python_sdk.constants.instrumentation.pinecone import APIS
+import json
 
 
 def generic_patch(original_method, method, version, tracer):
@@ -49,12 +52,21 @@ def generic_patch(original_method, method, version, tracer):
         attributes = DatabaseSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(api["METHOD"], kind=SpanKind.CLIENT) as span:
+
+            if span.is_recording():
+                set_span_attribute(span, "server.address", instance._config.host)
+                if method == "QUERY":
+                    set_query_input_attributes(span, kwargs)
+
             for field, value in attributes.model_dump(by_alias=True).items():
                 if value is not None:
                     span.set_attribute(field, value)
             try:
                 # Attempt to call the original method
                 result = original_method(instance, *args, **kwargs)
+                if result:
+                    if span.is_recording():
+                        set_query_response_attributes(span, result)
                 span.set_status(StatusCode.OK)
                 return result
             except Exception as err:
@@ -67,4 +79,48 @@ def generic_patch(original_method, method, version, tracer):
                 # Reraise the exception to ensure it's not swallowed
                 raise
 
+    @silently_fail
+    def set_query_input_attributes(span, kwargs):
+        set_span_attribute(span, "db.query.top_k", kwargs.get("top_k"))
+        set_span_attribute(span, "db.query.namespace", kwargs.get("namespace"))
+        set_span_attribute(span, "db.query.id", kwargs.get("id"))
+        filter = (
+            json.dumps(kwargs.get("filter"))
+            if isinstance(kwargs.get("filter"), dict)
+            else kwargs.get("filter")
+        )
+        set_span_attribute(span, "db.query.filter", filter)
+        set_span_attribute(
+            span, "db.query.include_values", kwargs.get("include_values")
+        )
+        set_span_attribute(
+            span, "db.query.include_metadata", kwargs.get("include_metadata")
+        )
+
+    @silently_fail
+    def set_query_response_attributes(span, response):
+        matches = response.get("matches")
+
+        usage = response.get("usage")
+        for match in matches:
+            span.add_event(
+                name="db.query.match",
+                attributes={
+                    "db.query.match.id": match.get("id"),
+                    "db.query.match.score": match.get("score"),
+                    "db.query.match.metadata": match.get("metadata"),
+                    # "db.query.match.values": match.get("values"),
+                },
+            )
+
+        if "read_units" in usage:
+            set_span_attribute(
+                span, "db.query.usage.read_units", usage.get("read_units")
+            )
+
+        if "write_units" in usage:
+            set_span_attribute(
+                span, "db.query.usage.write_units", usage.get("write_units")
+            )
+
     return traced_method

langtrace_python_sdk/instrumentation/qdrant/patch.py

@@ -15,6 +15,8 @@ limitations under the License.
 """
 
 from langtrace.trace_attributes import DatabaseSpanAttributes
+from langtrace_python_sdk.utils.silently_fail import silently_fail
+from langtrace_python_sdk.utils.llm import set_span_attributes
 from opentelemetry import baggage
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
@@ -47,12 +49,19 @@ def collection_patch(method, version, tracer):
             **(extra_attributes if extra_attributes is not None else {}),
         }
 
-        if hasattr(instance, "name") and instance.name is not None:
-            span_attributes["db.collection.name"] = instance.name
-
         attributes = DatabaseSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(api["METHOD"], kind=SpanKind.CLIENT) as span:
+            collection_name = kwargs.get("collection_name") or args[0]
+            operation = api["OPERATION"]
+            set_span_attributes(span, "db.collection.name", collection_name)
+            if operation == "upsert":
+                _set_upsert_attributes(span, args, kwargs)
+            elif operation == "add":
+                _set_upload_attributes(span, args, kwargs, "documents")
+
+            # Todo: Add support for other operations here.
+
             for field, value in attributes.model_dump(by_alias=True).items():
                 if value is not None:
                     span.set_attribute(field, value)
@@ -72,3 +81,26 @@ def collection_patch(method, version, tracer):
                 raise
 
     return traced_method
+
+
+@silently_fail
+def _set_upsert_attributes(span, args, kwargs):
+    points = kwargs.get("points") or args[1]
+    if isinstance(points, list):
+        length = len(points)
+    else:
+        # In case of using Batch.
+        length = len(points.ids)
+    set_span_attributes(span, "db.upsert.points_count", length)
+
+
+@silently_fail
+def _set_upload_attributes(span, args, kwargs, field):
+    docs = kwargs.get(field) or args[0]
+    if isinstance(docs, list):
+        length = len(docs)
+    else:
+        # In case of using Batch.
+        length = len(docs.ids)
+
+    set_span_attributes(span, f"db.upload.{field}_count", length)

langtrace_python_sdk/langtrace.py

@@ -24,6 +24,8 @@ from opentelemetry.sdk.trace.export import (
     ConsoleSpanExporter,
     SimpleSpanProcessor,
 )
+import sys
+from opentelemetry.sdk.resources import Resource
 
 from langtrace_python_sdk.extensions.langtrace_exporter import LangTraceExporter
 from langtrace_python_sdk.instrumentation.anthropic.instrumentation import (
@@ -72,7 +74,7 @@ def init(
     api_host: Optional[str] = None,
     disable_instrumentations: Optional[DisableInstrumentations] = None,
 ):
-    provider = TracerProvider()
+    provider = TracerProvider(resource=Resource.create({"service.name": sys.argv[0]}))
 
     remote_write_exporter = (
         LangTraceExporter(api_key=api_key, api_host=api_host)
@@ -128,24 +130,29 @@ def init_instrumentations(
     if disable_instrumentations is None:
         for _, v in all_instrumentations.items():
             v.instrument()
-            return
-    validate_instrumentations(disable_instrumentations)
-
-    for key in disable_instrumentations:
-        for vendor in disable_instrumentations[key]:
-            if key == "only":
-                filtered_dict = {
-                    k: v for k, v in all_instrumentations.items() if k != vendor.value
-                }
-                for _, v in filtered_dict.items():
-                    v.instrument()
-            else:
-                filtered_dict = {
-                    k: v for k, v in all_instrumentations.items() if k == vendor.value
-                }
-
-                for _, v in filtered_dict.items():
-                    v.instrument()
+    else:
+
+        validate_instrumentations(disable_instrumentations)
+
+        for key in disable_instrumentations:
+            for vendor in disable_instrumentations[key]:
+                if key == "only":
+                    filtered_dict = {
+                        k: v
+                        for k, v in all_instrumentations.items()
+                        if k != vendor.value
+                    }
+                    for _, v in filtered_dict.items():
+                        v.instrument()
+                else:
+                    filtered_dict = {
+                        k: v
+                        for k, v in all_instrumentations.items()
+                        if k == vendor.value
+                    }
+
+                    for _, v in filtered_dict.items():
+                        v.instrument()
 
 
 def validate_instrumentations(disable_instrumentations):

langtrace_python_sdk/utils/__init__.py

@@ -0,0 +1,7 @@
+def set_span_attribute(span, name, value):
+    if value is not None:
+        if value != "":
+            span.set_attribute(name, value)
+    return
+
+

langtrace_python_sdk/utils/llm.py

@@ -58,3 +58,11 @@ def calculate_price_from_usage(model, usage):
             + cost_table["output"] * usage["completion_tokens"]
         ) / 1000
     return 0
+
+
+def set_span_attributes(span, name, value):
+    if value is not None:
+        if value != "":
+            span.set_attribute(name, value)
+    return
+

langtrace_python_sdk/utils/prompt_registry.py

@@ -0,0 +1,82 @@
+import os
+from langtrace_python_sdk.constants.exporter.langtrace_exporter import (
+    LANGTRACE_REMOTE_URL,
+)
+import requests
+from urllib.parse import urlencode
+from typing import Optional, TypedDict, Dict, List
+
+
+class LangtracePrompt(TypedDict):
+    id: str
+    value: str
+    variables: List[str]
+    model: str
+    modelSettings: Dict[str, object]
+    version: int
+    live: bool
+    tags: List[str]
+    note: str
+    promptsetId: str
+    createdAt: str
+    updatedAt: str
+
+
+class FetchOptions(TypedDict, total=False):
+    prompt_version: int
+    variables: Dict[str, str]
+
+
+def get_prompt_from_registry(
+    prompt_registry_id: str,
+    options: Optional[FetchOptions] = None,
+) -> LangtracePrompt:
+    """Fetches a prompt from the registry.
+
+    Args:
+        prompt_registry_id (str): The ID of the prompt registry.
+        options (dict, optional): Configuration options for fetching the prompt:
+            - prompt_version (int): Fetches the prompt with the specified version.
+            - variables (dict): Replaces variables in the prompt with provided values.
+
+    Returns:
+        dict: The fetched prompt with variables replaced as specified.
+
+    Raises:
+        Exception: If the fetch operation fails or returns an error.
+    """
+    try:
+        query_params = {"promptset_id": prompt_registry_id}
+        if options:
+            if "prompt_version" in options:
+                query_params["version"] = options["prompt_version"]
+            if "variables" in options:
+                for key, value in options["variables"].items():
+                    query_params[f"variables.{key}"] = value
+        # Encode the query parameters
+        query_string = urlencode(query_params, doseq=True)
+        headers = {"x-api-key": os.environ["LANGTRACE_API_KEY"]}
+        print(query_params)
+        # Make the GET request to the API
+        response = requests.get(
+            f"{LANGTRACE_REMOTE_URL}/api/promptset?{query_string}",
+            headers=headers,
+            timeout=None,
+        )
+        response.raise_for_status()
+
+        # Extract the prompt data from the response
+        prompt_data = response.json()["prompts"][0]
+        return prompt_data
+
+    except requests.RequestException as err:
+        # Handle specific HTTP errors or general request exceptions
+        error_msg = str(err)
+        if err.response:
+            try:
+                # Try to extract server-provided error message
+                error_msg = err.response.json().get("error", error_msg)
+            except ValueError:
+                # Fallback if response is not JSON
+                error_msg = err.response.text
+        raise Exception(f"API error: {error_msg}") from err

langtrace_python_sdk/utils/silently_fail.py

@@ -0,0 +1,19 @@
+import logging
+
+
+def silently_fail(func):
+    """
+    A decorator that catches exceptions thrown by the decorated function and logs them as warnings.
+    """
+
+    logger = logging.getLogger(func.__module__)
+
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as exception:
+            logger.warning(
+                "Failed to execute %s, error: %s", func.__name__, str(exception)
+            )
+
+    return wrapper

langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "2.0.6"
+__version__ = "2.0.8"

{langtrace_python_sdk-2.0.6.dist-info → langtrace_python_sdk-2.0.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: langtrace-python-sdk
-Version: 2.0.6
+Version: 2.0.8
 Summary: Python SDK for LangTrace
 Project-URL: Homepage, https://github.com/Scale3-Labs/langtrace-python-sdk
 Author-email: Scale3 Labs <engineering@scale3labs.com>
@@ -25,7 +25,9 @@ Requires-Dist: langchain-openai; extra == 'dev'
 Requires-Dist: llama-index; extra == 'dev'
 Requires-Dist: openai; extra == 'dev'
 Requires-Dist: python-dotenv; extra == 'dev'
+Requires-Dist: qdrant-client; extra == 'dev'
 Provides-Extra: test
+Requires-Dist: protobuf==3.20.0; extra == 'test'
 Requires-Dist: pytest; extra == 'test'
 Requires-Dist: pytest-asyncio; extra == 'test'
 Requires-Dist: pytest-vcr; extra == 'test'
@@ -217,6 +219,15 @@ def chat_completion():
    api_call2()
 ```
 
+- `get_prompt_from_registry` - this function is designed to fetch the desired prompt from the `Prompt Registry`. You can pass two options for filtering `prompt_version` & `variables`.
+
+
+```python
+from langtrace_python_sdk import get_prompt_from_registry
+
+prompt = get_prompt_from_registry(<Registry ID>, options={"prompt_version": 1, "variables": {"foo": "bar"} })
+```
+
 ## Supported integrations
 
 Langtrace automatically captures traces from the following vendors:

{langtrace_python_sdk-2.0.6.dist-info → langtrace_python_sdk-2.0.8.dist-info}/RECORD

@@ -32,12 +32,12 @@ examples/openai_example/tool_calling_nonstreaming.py,sha256=Yc848IooZRXNynHL6z0k
 examples/openai_example/tool_calling_streaming.py,sha256=mV1RbyAoVhumGRPpqPWQ6PMhnJyeifrlELd2-K1qJ_w,7015
 examples/perplexity_example/basic.py,sha256=bp7n27gaugJkaFVyt8pjaEfi66lYcqP6eFFjPewUShY,668
 examples/pinecone_example/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-examples/pinecone_example/basic.py,sha256=9IWMoMLozMTpxjY8BFBg2EpUNY5dDjU1xocEG9XsXP0,966
+examples/pinecone_example/basic.py,sha256=uwWu_UFuDtthgs-x8gOPnsUGOeBDCVPEoiz9uwSEjfg,1007
 examples/qdrant_example/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 examples/qdrant_example/basic.py,sha256=DCMjHSuBZKkhEjCkwy5d5La9WMyW0lCWqtcZWiFCEm4,1425
-langtrace_python_sdk/__init__.py,sha256=q2TAsxGIv4wFIQA7G-HGxlnzO89bhPU5TgI42sPHLj0,740
-langtrace_python_sdk/langtrace.py,sha256=z6CL_e2aQUXJlTUzel0IL2zBVAa5lxRR01b88GfaBuY,6576
-langtrace_python_sdk/version.py,sha256=4Yxft_MNOtwQG58mZpa8eoM0XBHa8VNkTPH9_i0gDaM,22
+langtrace_python_sdk/__init__.py,sha256=_r2S-lLZFTFFvWjf0lWBX4nzkJN5R7FHX0Dnmr76CEk,848
+langtrace_python_sdk/langtrace.py,sha256=o0xJInE_dwNBkAfIQRqMKPz6glCAcGdRJxEG8kasq8I,6839
+langtrace_python_sdk/version.py,sha256=oVbl6ILP5dyBZtm2iUzM_QTUhSy3C325aGgDThF7P_4,22
 langtrace_python_sdk/constants/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langtrace_python_sdk/constants/exporter/langtrace_exporter.py,sha256=5MNjnAOg-4am78J3gVMH6FSwq5N8TOj72ugkhsw4vi0,46
 langtrace_python_sdk/constants/instrumentation/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -84,13 +84,16 @@ langtrace_python_sdk/instrumentation/openai/instrumentation.py,sha256=G2HSZfr6Du
 langtrace_python_sdk/instrumentation/openai/patch.py,sha256=kQxc_E1IlOgHZnlc0HbpGjM5rZCDR3YpNFChp2X9GbM,37263
 langtrace_python_sdk/instrumentation/pinecone/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langtrace_python_sdk/instrumentation/pinecone/instrumentation.py,sha256=52GbkfRK-sxXBXhHmHRRO2BGNd6zZVHDp8njLS1CQX0,2314
-langtrace_python_sdk/instrumentation/pinecone/patch.py,sha256=ebz57WgmlTnY-FDwQgJ-_eqsknEmDWfh-AiNow4G94s,2705
+langtrace_python_sdk/instrumentation/pinecone/patch.py,sha256=hLgYWDKTytuvTORDd6dupMpKAp19vCgOAXDOdjvWme0,4832
 langtrace_python_sdk/instrumentation/qdrant/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langtrace_python_sdk/instrumentation/qdrant/instrumentation.py,sha256=vl2eKSP55aqDo1JiRlvOUBrr6kddvG9Z5dCYew2OG08,1816
-langtrace_python_sdk/instrumentation/qdrant/patch.py,sha256=skVziJbvoR7XYTI6KmsiQlclTmQEeLeJ5UrJLvi7fjw,2807
+langtrace_python_sdk/instrumentation/qdrant/patch.py,sha256=x95PXaJe_GrOG51zZCDuchVA-bzuA7s5LHqOoWrqSA0,3883
 langtrace_python_sdk/types/__init__.py,sha256=QYK-kRkyfXDdzjjKmXVxi0vZ6Wpbxdf0CoZbTXO0JT8,787
-langtrace_python_sdk/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-langtrace_python_sdk/utils/llm.py,sha256=Cl710-0rfr15g1xNkOpLpgCn84pNlfMflcQ4S5jSRu0,2010
+langtrace_python_sdk/utils/__init__.py,sha256=is6gOLiHcvrN9kPgYnUwyzkgz9vInVQ4jGe7k0zRHrA,150
+langtrace_python_sdk/utils/llm.py,sha256=gnmJ6YWzzdOCaVFg1C-iyNcMGDyiJ623qZeOwDYUfIQ,2162
+langtrace_python_sdk/utils/prompt_registry.py,sha256=7FFB4Pj0414qgf02h5zL5vXBZgNBf74g4Iq7GdFaIO0,2689
+langtrace_python_sdk/utils/silently_fail.py,sha256=F_9EteXCO9Cyq-8MA1OT2Zy_dx8n06nt31I7t7ui24E,478
+langtrace_python_sdk/utils/types.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langtrace_python_sdk/utils/with_root_span.py,sha256=MXF-FmG47P0GtJQNffKM_3XGrJHVAC4hMB6GmwtnA4M,2454
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/conftest.py,sha256=0Jo6iCZTXbdvyJVhG9UpYGkLabL75378oauCzmt-Sa8,603
@@ -100,7 +103,8 @@ tests/anthropic/test_anthropic.py,sha256=p4WhcECJ_zctirtkAX6WejrlNfZzo90fSSxqQki
 tests/anthropic/cassettes/test_anthropic.yaml,sha256=z7YAqA_BBgI-hw7uyVMdLoIZdBYhKwl9cuTzUu9nAZs,2332
 tests/anthropic/cassettes/test_anthropic_streaming.yaml,sha256=D0w4-6dfsgrhbNLJEj8gZBV0yXfrAfA9u90Yu8Ac-TY,11675
 tests/anthropic/cassettes/test_async_anthropic_streaming.yaml,sha256=hQZPY2vwBaW3BWllLd0lcGQ73DjA8C3Ips1Hx9pA-ao,8373
-tests/chroma/test_chroma.py,sha256=bVWC_OwmarTgVzKxZy5lrXpJzu9v9Ko6uQdktxOVh74,2430
+tests/chroma/conftest.py,sha256=kqb4VydCnlVpkvBX5Bu-pfnVS-ZZfju9cp6vq76tkJI,309
+tests/chroma/test_chroma.py,sha256=-KJHunvvVi1OoKKOcKCeHO1s399Gm9vJfd-EzgllQmk,1220
 tests/cohere/conftest.py,sha256=jBbyg-tut-ZJN5_5D64sGmaPIhT_nQQQAiW43kl5Rz4,621
 tests/cohere/test_cohere_chat.py,sha256=MioEIXB5e-4t99Vt-KpLlxd4diPbp3S6eNfnPWjEDXM,4564
 tests/cohere/test_cohere_embed.py,sha256=bFKDzUnypMyLefAOHR-6dEgvt0yys2Yw7-F0zc0uQyc,1250
@@ -121,8 +125,13 @@ tests/openai/cassettes/test_async_image_generation.yaml,sha256=_LYZcrqxrnSqcWVQn
 tests/openai/cassettes/test_chat_completion.yaml,sha256=YkNFgK9VHAzNqGWuxFcTiE194GdEie8eDf1FSsffjd8,2944
 tests/openai/cassettes/test_chat_completion_streaming.yaml,sha256=nkx_TemQMYSZxUF_b-LCEFwCRDm0AkQHLf4sdJVuZBw,2592394
 tests/openai/cassettes/test_image_generation.yaml,sha256=gn5aSVp6V6_hb_rt2NnkAWd_idzDxo-7VzhZII0Wslw,3562
-tests/pinecone/test_pinecone.py,sha256=iFuAqzbfrxHzH_J2kDCbezqYv1m-yeQiHiO3rMKTuYc,2735
-langtrace_python_sdk-2.0.6.dist-info/METADATA,sha256=To2Jd-L4FM7sHtVWbVbSU19s0BQKynZm9LKboKRLrmQ,9582
-langtrace_python_sdk-2.0.6.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
-langtrace_python_sdk-2.0.6.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-langtrace_python_sdk-2.0.6.dist-info/RECORD,,
+tests/pinecone/conftest.py,sha256=EY1m5M6MKSOktbVEe__wYNjCW9A7H0xdC-aXhU8JEOg,916
+tests/pinecone/test_pinecone.py,sha256=JMe3S15cfBWjTArC-FTCarRRYBFylkv4BwQgL2BmUfw,3614
+tests/pinecone/cassettes/test_query.yaml,sha256=b5v9G3ssUy00oG63PlFUR3JErF2Js-5AIqHnMsvJM8w,103821
+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-2.0.8.dist-info/METADATA,sha256=bUB2J9h4OXlPWUXjo-VhX2bqamXaImnC9_vbdxiH5Uw,10046
+langtrace_python_sdk-2.0.8.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
+langtrace_python_sdk-2.0.8.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+langtrace_python_sdk-2.0.8.dist-info/RECORD,,

tests/chroma/conftest.py

@@ -0,0 +1,15 @@
+import chromadb
+import pytest
+from langtrace_python_sdk.instrumentation.chroma.instrumentation import (
+    ChromaInstrumentation,
+)
+
+
+@pytest.fixture
+def chroma_client():
+    return chromadb.Client()
+
+
+@pytest.fixture(scope="session", autouse=True)
+def instrument():
+    ChromaInstrumentation().instrument()

tests/chroma/test_chroma.py

@@ -1,70 +1,31 @@
-from unittest.mock import MagicMock, patch, call
+import importlib
+from chromadb.utils import embedding_functions
 from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
-from langtrace_python_sdk.constants.instrumentation.openai import APIS
-from opentelemetry.trace.status import StatusCode, Status
-from langtrace_python_sdk.instrumentation.chroma.patch import collection_patch
-from opentelemetry.trace import SpanKind
-from tests.utils import common_setup
-import unittest
-import json
 
-
-class TestChromaPatch(unittest.TestCase):
-    data = {
-        "status": "success",
-    }
-
-    def setUp(self):
-        self.chroma_mock, self.tracer, self.span = common_setup(
-            self.data, "chromadb.Collection.add"
-        )
-        self.wrapped_method = MagicMock(return_value="mocked method result")
-        self.instance = MagicMock()
-        self.instance.name = "aa"
-
-    def tearDown(self):
-        self.chroma_mock.stop()
-
-    def test_collection_patch_success(self):
-        # Arrange
-        traced_method = collection_patch("ADD", "1.2.3", self.tracer)
-
-        # Act
-        result = traced_method(self.wrapped_method, self.instance, (), {})
-
-        # Assert
-        # Assert the result of the original method is returned
-        self.assertEqual(result, "mocked method result")
-
-        # Assert the span is started with the correct parameters
-        self.assertTrue(
-            self.tracer.start_as_current_span.called_once_with(
-                "chromadb.Collection.add", kind=SpanKind.CLIENT
-            )
-        )
-
-        # Verify span attributes are set as expected
-        expected_attributes = {
-            "langtrace.sdk.name": "langtrace-python-sdk",
-            "langtrace.service.name": "Chroma",
-            "langtrace.service.type": "vectordb",
-            "langtrace.service.version": "1.2.3",
-            "langtrace.version": "1.0.0",
-            "db.system": "chromadb",
-            "db.operation": "add",
-            "db.collection.name": "aa",
-        }
-        for key, value in expected_attributes.items():
-            self.span.set_attribute.assert_has_calls([call(key, value)], any_order=True)
-
-        actual_calls = self.span.set_attribute.call_args_list
-
-        for key, value in expected_attributes.items():
-            self.assertIn(call(key, value), actual_calls)
-
-        # Assert the span status is set to OK
-        self.span.set_status.assert_called_with(StatusCode.OK)
-
-
-if __name__ == "__main__":
-    unittest.main()
+COLLECTION_NAME = "test_collection"
+
+
+def test_chroma_add(chroma_client, exporter):
+    embedder = embedding_functions.DefaultEmbeddingFunction()
+    collection = chroma_client.create_collection(
+        name=COLLECTION_NAME, embedding_function=embedder
+    )
+    collection.add(
+        documents=["This is a document", "This is another document"],
+        metadatas=[{"source": "my_source"}, {"source": "my_source"}],
+        ids=["id1", "id2"],
+    )
+    spans = exporter.get_finished_spans()
+    chroma_span = spans[-1]
+
+    attributes = chroma_span.attributes
+
+    assert attributes.get("db.collection.name") == COLLECTION_NAME
+    assert attributes.get("db.operation") == "add"
+    assert attributes.get("db.system") == "chromadb"
+    assert attributes.get("langtrace.sdk.name") == "langtrace-python-sdk"
+    assert attributes.get("langtrace.service.name") == SERVICE_PROVIDERS["CHROMA"]
+    assert attributes.get("langtrace.service.type") == "vectordb"
+    assert attributes.get("langtrace.service.version") == importlib.metadata.version(
+        "chromadb"
+    )