langtrace-python-sdk 2.0.13__py3-none-any.whl → 2.1.1__py3-none-any.whl

langtrace_python_sdk/__init__.py

@@ -17,5 +17,11 @@ 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
+from langtrace_python_sdk.utils.with_root_span import send_user_feedback
 
-__all__ = ["langtrace", "with_langtrace_root_span", "get_prompt_from_registry"]
+__all__ = [
+    "langtrace",
+    "with_langtrace_root_span",
+    "get_prompt_from_registry",
+    "send_user_feedback",
+]

langtrace_python_sdk/constants/instrumentation/common.py

@@ -22,6 +22,7 @@ SERVICE_PROVIDERS = {
     "COHERE": "Cohere",
     "PPLX": "Perplexity",
     "QDRANT": "Qdrant",
+    "WEAVIATE": "Weaviate",
 }
 
 LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY = "langtrace_additional_attributes"

langtrace_python_sdk/constants/instrumentation/weaviate.py

@@ -0,0 +1,44 @@
+from langtrace.trace_attributes import WeaviateMethods
+
+APIS = {
+    "weaviate.query.bm25": {
+        "MODULE": WeaviateMethods.QUERY_BM25.value,
+        "METHOD": "_BM25Query.bm25",
+        "OPERATION": "query",
+    },
+    "weaviate.query.fetch_object_by_id": {
+        "MODULE": WeaviateMethods.QUERY_FETCH_OBJECT_BY_ID.value,
+        "METHOD": "_FetchObjectByIDQuery.fetch_object_by_id",
+        "OPERATION": "query",
+    },
+    "weaviate.query.fetch_objects": {
+        "MODULE": WeaviateMethods.QUERY_FETCH_OBJECTS.value,
+        "METHOD": "_FetchObjectsQuery.fetch_objects",
+        "OPERATION": "query",
+    },
+    "weaviate.query.hybrid": {
+        "MODULE": WeaviateMethods.QUERY_HYBRID.value,
+        "METHOD": "_HybridQuery.hybrid",
+        "OPERATION": "query",
+    },
+    "weaviate.query.near_object": {
+        "MODULE": WeaviateMethods.QUERY_NEAR_OBJECT.value,
+        "METHOD": "_NearObjectQuery.near_object",
+        "OPERATION": "query",
+    },
+    "weaviate.query.near_text": {
+        "MODULE": WeaviateMethods.QUERY_NEAR_TEXT.value,
+        "METHOD": "_NearTextQuery.near_text",
+        "OPERATION": "query",
+    },
+    "weaviate.query.near_vector": {
+        "MODULE": WeaviateMethods.QUERY_NEAR_VECTOR.value,
+        "METHOD": "_NearVectorQuery.near_vector",
+        "OPERATION": "query",
+    },
+    "weaviate.collections.create": {
+        "MODULE": WeaviateMethods.COLLECTIONS_OPERATIONS.value,
+        "METHOD": "_Collections.create",
+        "OPERATION": "create",
+    },
+}

langtrace_python_sdk/instrumentation/chroma/patch.py

@@ -15,6 +15,8 @@ limitations under the License.
 """
 
 from langtrace.trace_attributes import DatabaseSpanAttributes
+from langtrace_python_sdk.utils.llm import set_span_attributes
+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 (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
     SERVICE_PROVIDERS,
 )
+import json
 
 
 def collection_patch(method, version, tracer):
@@ -44,6 +47,7 @@ def collection_patch(method, version, tracer):
             "langtrace.version": "1.0.0",
             "db.system": "chromadb",
             "db.operation": api["OPERATION"],
+            "db.query": json.dumps(kwargs.get("query")),
             **(extra_attributes if extra_attributes is not None else {}),
         }
 
@@ -57,8 +61,32 @@ def collection_patch(method, version, tracer):
                 if value is not None:
                     span.set_attribute(field, value)
             try:
+
+                operation = api["OPERATION"]
+                if operation == "add":
+                    _set_chroma_add_attributes(span, kwargs)
+                elif operation == "get":
+                    _set_chroma_get_attributes(span, kwargs)
+                elif operation == "query":
+                    _set_chroma_query_attributes(span, kwargs)
+                elif operation == "peek":
+                    _set_chroma_peek_attributes(span, kwargs)
+                elif operation == "update":
+                    _set_chroma_update_attributes(span, kwargs)
+                elif operation == "upsert":
+                    _set_chroma_upsert_attributes(span, kwargs)
+                elif operation == "modify":
+                    _set_chroma_modify_attributes(span, kwargs)
+                elif operation == "delete":
+                    _set_chroma_delete_attributes(span, kwargs)
                 # Attempt to call the original method
                 result = wrapped(*args, **kwargs)
+
+                if operation == "query":
+                    events = _set_chroma_query_response(span, result)
+                    for event in events:
+                        span.add_event(name="db.chroma.query.result", attributes=event)
+
                 span.set_status(StatusCode.OK)
                 return result
             except Exception as err:
@@ -72,3 +100,166 @@ def collection_patch(method, version, tracer):
                 raise
 
     return traced_method
+
+
+def get_count_or_none(value):
+    return len(value) if value is not None else None
+
+
+def handle_null_params(param):
+    return str(param) if param else None
+
+
+@silently_fail
+def _set_chroma_add_attributes(span, kwargs):
+    set_span_attributes(
+        span, "db.chroma.add.ids_count", get_count_or_none(kwargs.get("ids"))
+    )
+    set_span_attributes(
+        span,
+        "db.chroma.add.embeddings_count",
+        get_count_or_none(kwargs.get("embeddings")),
+    )
+    set_span_attributes(
+        span,
+        "db.chroma.add.metadatas_count",
+        get_count_or_none(kwargs.get("metadatas")),
+    )
+    set_span_attributes(
+        span,
+        "db.chroma.add.documents_count",
+        get_count_or_none(kwargs.get("documents")),
+    )
+
+
+@silently_fail
+def _set_chroma_get_attributes(span, kwargs):
+    set_span_attributes(
+        span, "db.chroma.get.ids_count", get_count_or_none(kwargs.get("ids"))
+    )
+    set_span_attributes(
+        span, "db.chroma.get.where", handle_null_params(kwargs.get("where"))
+    )
+    set_span_attributes(span, "db.chroma.get.limit", kwargs.get("limit"))
+    set_span_attributes(span, "db.chroma.get.offset", kwargs.get("offset"))
+    set_span_attributes(
+        span,
+        "db.chroma.get.where_document",
+        handle_null_params(kwargs.get("where_document")),
+    )
+    set_span_attributes(
+        span, "db.chroma.get.include", handle_null_params(kwargs.get("include"))
+    )
+
+
+@silently_fail
+def _set_chroma_query_attributes(span, kwargs):
+    set_span_attributes(
+        span,
+        "db.chroma.query.query_embeddings_count",
+        get_count_or_none(kwargs.get("query_embeddings")),
+    )
+    set_span_attributes(
+        span,
+        "db.chroma.query.query_texts_count",
+        get_count_or_none(kwargs.get("query_texts")),
+    )
+    set_span_attributes(span, "db.chroma.query.n_results", kwargs.get("n_results"))
+    set_span_attributes(
+        span, "db.chroma.query.where", handle_null_params(kwargs.get("where"))
+    )
+    set_span_attributes(
+        span,
+        "db.chroma.query.where_document",
+        handle_null_params(kwargs.get("where_document")),
+    )
+    set_span_attributes(
+        span, "db.chroma.query.include", handle_null_params(kwargs.get("include"))
+    )
+
+
+@silently_fail
+def _set_chroma_peek_attributes(span, kwargs):
+    set_span_attributes(span, "db.chroma.peek.limit", kwargs.get("limit"))
+
+
+@silently_fail
+def _set_chroma_update_attributes(span, kwargs):
+    set_span_attributes(
+        span, "db.chroma.update.ids_count", get_count_or_none(kwargs.get("ids"))
+    )
+    set_span_attributes(
+        span,
+        "db.chroma.update.embeddings_count",
+        get_count_or_none(kwargs.get("embeddings")),
+    )
+    set_span_attributes(
+        span,
+        "db.chroma.update.metadatas_count",
+        get_count_or_none(kwargs.get("metadatas")),
+    )
+    set_span_attributes(
+        span,
+        "db.chroma.update.documents_count",
+        get_count_or_none(kwargs.get("documents")),
+    )
+
+
+@silently_fail
+def _set_chroma_modify_attributes(span, kwargs):
+    set_span_attributes(span, "db.chroma.modify.name", kwargs.get("name"))
+    # TODO: Add metadata attribute
+
+
+@silently_fail
+def _set_chroma_upsert_attributes(span, kwargs):
+    set_span_attributes(
+        span,
+        "db.chroma.upsert.embeddings_count",
+        get_count_or_none(kwargs.get("embeddings")),
+    )
+    set_span_attributes(
+        span,
+        "db.chroma.upsert.metadatas_count",
+        get_count_or_none(kwargs.get("metadatas")),
+    )
+    set_span_attributes(
+        span,
+        "db.chroma.upsert.documents_count",
+        get_count_or_none(kwargs.get("documents")),
+    )
+
+
+@silently_fail
+def _set_chroma_delete_attributes(span, kwargs):
+    set_span_attributes(
+        span, "db.chroma.delete.ids_count", get_count_or_none(kwargs.get("ids"))
+    )
+    set_span_attributes(
+        span, "db.chroma.delete.where", handle_null_params(kwargs.get("where"))
+    )
+    set_span_attributes(
+        span,
+        "db.chroma.delete.where_document",
+        handle_null_params(kwargs.get("where_document")),
+    )
+
+
+@silently_fail
+def _set_chroma_query_response(span, result):
+
+    attributes = []
+    ids = result.get("ids")[0]
+    distances = result.get("distances")[0]
+    metadatas = result.get("metadatas")[0]
+    documents = result.get("documents")[0]
+
+    for idx, _ in enumerate(ids):
+        attribute = {
+            "id": ids[idx],
+            "distance": distances[idx],
+            "metadata": metadatas[idx],
+            "document": documents[idx],
+        }
+        attributes.append(attribute)
+    return attributes

langtrace_python_sdk/instrumentation/openai/patch.py

@@ -616,6 +616,7 @@ def async_chat_completions_create(original_method, version, tracer):
         span.add_event(Event.STREAM_START.value)
         completion_tokens = 0
         try:
+            content = []
             async for chunk in result:
                 if hasattr(chunk, "model") and chunk.model is not None:
                     span.set_attribute("llm.model", chunk.model)

langtrace_python_sdk/instrumentation/pinecone/patch.py

@@ -46,6 +46,7 @@ def generic_patch(original_method, method, version, tracer):
             "langtrace.version": "1.0.0",
             "db.system": "pinecone",
             "db.operation": api["OPERATION"],
+            "db.query": json.dumps(kwargs.get("query")),
             **(extra_attributes if extra_attributes is not None else {}),
         }
 

langtrace_python_sdk/instrumentation/qdrant/patch.py

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import json
 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
@@ -46,6 +47,7 @@ def collection_patch(method, version, tracer):
             "langtrace.version": "1.0.0",
             "db.system": "qdrant",
             "db.operation": api["OPERATION"],
+            "db.query": json.dumps(kwargs.get("query")),
             **(extra_attributes if extra_attributes is not None else {}),
         }
 
@@ -55,12 +57,22 @@ def collection_patch(method, version, tracer):
             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":
+
+            if operation == "add":
                 _set_upload_attributes(span, args, kwargs, "documents")
 
-            # Todo: Add support for other operations here.
+            elif operation == "upsert":
+                _set_upsert_attributes(span, args, kwargs)
+
+            elif operation in ["query", "discover", "recommend", "retrieve", "search"]:
+                _set_search_attributes(span, args, kwargs)
+            elif operation in [
+                "query_batch",
+                "discover_batch",
+                "recommend_batch",
+                "search_batch",
+            ]:
+                _set_batch_search_attributes(span, args, kwargs, operation)
 
             for field, value in attributes.model_dump(by_alias=True).items():
                 if value is not None:
@@ -104,3 +116,15 @@ def _set_upload_attributes(span, args, kwargs, field):
         length = len(docs.ids)
 
     set_span_attributes(span, f"db.upload.{field}_count", length)
+
+
+@silently_fail
+def _set_search_attributes(span, args, kwargs):
+    limit = kwargs.get("limit") or 10
+    set_span_attributes(span, "db.query.top_k", limit)
+
+
+@silently_fail
+def _set_batch_search_attributes(span, args, kwargs, method):
+    requests = kwargs.get("requests") or []
+    set_span_attributes(span, f"db.{method}.requests_count", len(requests))

langtrace_python_sdk/instrumentation/weaviate/instrumentation.py

@@ -0,0 +1,66 @@
+"""
+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.
+"""
+
+import importlib.metadata
+import logging
+from typing import Collection
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from langtrace_python_sdk.instrumentation.weaviate.patch import (
+    generic_collection_patch,
+    generic_query_patch,
+)
+from langtrace_python_sdk.constants.instrumentation.weaviate import APIS
+
+logging.basicConfig(level=logging.FATAL)
+
+
+class WeaviateInstrumentation(BaseInstrumentor):
+    """
+    The WeaviateInstrumentation class represents the instrumentation for the Weaviate SDK.
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["weaviate-client >= 4.6.1", "trace-attributes >= 4.0.2"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = importlib.metadata.version("weaviate-client")
+
+        for api_name, api_config in APIS.items():
+            module_path, function_name = api_name.rsplit(".", 1)
+            if api_config.get("OPERATION") == "query":
+                wrap_function_wrapper(
+                    api_config["MODULE"],
+                    api_config["METHOD"],
+                    generic_query_patch(api_name, version, tracer),
+                )
+            elif api_config.get("OPERATION") == "create":
+                wrap_function_wrapper(
+                    api_config["MODULE"],
+                    api_config["METHOD"],
+                    generic_collection_patch(api_name, version, tracer),
+                )
+
+    def _instrument_module(self, module_name):
+        pass
+
+    def _uninstrument(self, **kwargs):
+        pass

langtrace_python_sdk/instrumentation/weaviate/patch.py

@@ -0,0 +1,166 @@
+"""
+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.
+"""
+
+import json
+
+from langtrace.trace_attributes import DatabaseSpanAttributes
+from opentelemetry import baggage
+from opentelemetry.trace import SpanKind
+from opentelemetry.trace.status import Status, StatusCode
+
+from langtrace_python_sdk.constants.instrumentation.common import (
+    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
+    SERVICE_PROVIDERS,
+)
+from langtrace_python_sdk.constants.instrumentation.weaviate import APIS
+from langtrace_python_sdk.utils.misc import extract_input_params, to_iso_format
+
+# Predefined metadata response attributes
+METADATA_ATTRIBUTES = [
+    "creation_time",
+    "last_update_time",
+    "distance",
+    "certainty",
+    "score",
+    "explain_score",
+    "is_consistent",
+    "rerank_score",
+]
+
+
+def extract_metadata(metadata):
+    # Extraction response Query metadata
+    extracted_metadata = {
+        attr: (
+            to_iso_format(getattr(metadata, attr))
+            if "time" in attr
+            else getattr(metadata, attr)
+        )
+        for attr in METADATA_ATTRIBUTES
+        if hasattr(metadata, attr)
+    }
+
+    return {k: v for k, v in extracted_metadata.items() if v is not None}
+
+
+def aggregate_responses(result):
+    all_responses = []
+
+    if hasattr(result, "objects") and result.objects is not None:
+        for each_obj in result.objects:
+            # Loop for multiple object responses
+            response_attributes = get_response_object_attributes(each_obj)
+            all_responses.append(response_attributes)
+    else:
+        # For single object responses
+        all_responses = get_response_object_attributes(result)
+
+    return json.dumps(all_responses)
+
+
+def get_response_object_attributes(response_object):
+
+    response_attributes = {
+        **response_object.properties,
+        "uuid": str(response_object.uuid) if hasattr(response_object, "uuid") else None,
+        "collection": (
+            response_object.collection
+            if hasattr(response_object, "collection")
+            else None
+        ),
+        "vector": (
+            response_object.vector if hasattr(response_object, "vector") else None
+        ),
+        "references": (
+            response_object.references
+            if hasattr(response_object, "references")
+            else None
+        ),
+        "metadata": (
+            extract_metadata(response_object.metadata)
+            if hasattr(response_object, "metadata")
+            else None
+        ),
+    }
+    response_attributes = {
+        k: v for k, v in response_attributes.items() if v is not None
+    }
+    return response_attributes
+
+
+def create_traced_method(method_name, version, tracer, get_collection_name=None):
+    def traced_method(wrapped, instance, args, kwargs):
+        api = APIS[method_name]
+        service_provider = SERVICE_PROVIDERS["WEAVIATE"]
+        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
+
+        collection_name = (
+            get_collection_name(instance, kwargs)
+            if get_collection_name
+            else instance._name
+        )
+
+        span_attributes = {
+            "langtrace.sdk.name": "langtrace-python-sdk",
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "vectordb",
+            "langtrace.service.version": version,
+            "langtrace.version": "1.0.0",
+            "db.system": "weaviate",
+            "db.operation": api["OPERATION"],
+            "db.collection.name": collection_name,
+            "db.query": json.dumps(extract_input_params(args, kwargs)),
+            **(extra_attributes if extra_attributes is not None else {}),
+        }
+
+        attributes = DatabaseSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(method_name, kind=SpanKind.CLIENT) as span:
+            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 = wrapped(*args, **kwargs)
+                if api["OPERATION"] == "query":
+                    span.add_event(
+                        name="db.response",
+                        attributes={"db.response": aggregate_responses(result)},
+                    )
+                span.set_status(StatusCode.OK)
+                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 generic_query_patch(method_name, version, tracer):
+    return create_traced_method(method_name, version, tracer)
+
+
+def generic_collection_patch(method_name, version, tracer):
+    return create_traced_method(
+        method_name,
+        version,
+        tracer,
+        get_collection_name=lambda instance, kwargs: kwargs.get("name"),
+    )

langtrace_python_sdk/langtrace.py

@@ -64,6 +64,9 @@ from langtrace_python_sdk.instrumentation.pinecone.instrumentation import (
 from langtrace_python_sdk.instrumentation.qdrant.instrumentation import (
     QdrantInstrumentation,
 )
+from langtrace_python_sdk.instrumentation.weaviate.instrumentation import (
+    WeaviateInstrumentation,
+)
 
 
 def init(
@@ -119,6 +122,7 @@ def init(
         "langgraph": LanggraphInstrumentation(),
         "anthropic": AnthropicInstrumentation(),
         "cohere": CohereInstrumentation(),
+        "weaviate": WeaviateInstrumentation(),
     }
 
     init_instrumentations(disable_instrumentations, all_instrumentations)

langtrace_python_sdk/types/__init__.py

@@ -15,6 +15,7 @@ class InstrumentationType(Enum):
     LANGCHAIN_CORE = "langchain_core"
     LANGCHAIN_COMMUNITY = "langchain_community"
     LANGGRAPH = "langgraph"
+    WEAVIATE = "weaviate"
 
     @staticmethod
     def from_string(value: str):

langtrace_python_sdk/utils/__init__.py

@@ -3,5 +3,3 @@ def set_span_attribute(span, name, value):
         if value != "":
             span.set_attribute(name, value)
     return
-
-

langtrace_python_sdk/utils/llm.py

@@ -65,4 +65,3 @@ def set_span_attributes(span, name, value):
         if value != "":
             span.set_attribute(name, value)
     return
-

langtrace_python_sdk/utils/misc.py

@@ -0,0 +1,30 @@
+from datetime import datetime
+import json
+
+
+def extract_input_params(args, kwargs):
+    extracted_params = {}
+    for key, value in kwargs.items():
+        if hasattr(value, "__dict__"):
+            extracted_params[key] = json.dumps(vars(value))
+        else:
+            extracted_params[key] = value
+    for i, value in enumerate(args):
+        if hasattr(value, "__dict__"):
+            extracted_params[f"arg{i}"] = json.dumps(vars(value))
+        else:
+            extracted_params[f"arg{i}"] = value
+    # Remove None values
+    return {k: v for k, v in extracted_params.items() if v is not None}
+
+
+def to_iso_format(value):
+    return (
+        None
+        if value is None
+        else (
+            value.isoformat(timespec="microseconds") + "Z"
+            if isinstance(value, datetime)
+            else None
+        )
+    )