langtrace-python-sdk 1.0.9__py3-none-any.whl

instrumentation/langchain_core/patch.py

@@ -0,0 +1,168 @@
+"""
+This module contains the patching functions for the langchain_core package.
+"""
+import json
+
+from langtrace.trace_attributes import FrameworkSpanAttributes
+from opentelemetry.trace import SpanKind, StatusCode
+from opentelemetry.trace.status import Status
+
+from instrumentation.constants import SERVICE_PROVIDERS
+
+
+def generic_patch(method_name, task, tracer, version, trace_output=True, trace_input=True):
+    """
+    Wrapper function to trace a generic method.
+    method_name: The name of the method to trace.
+    task: The name used to identify the type of task in `generic_patch`.
+    tracer: The tracer object used in `generic_patch`.
+    version: The version parameter used in `generic_patch`.
+    trace_output: Whether to trace the output of the patched methods.
+    trace_input: Whether to trace the input of the patched methods.
+    """
+
+    def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS['LANGCHAIN']
+        span_attributes = {
+            'langtrace.service.name': service_provider,
+            'langtrace.service.type': 'framework',
+            'langtrace.service.version': version,
+            'langtrace.version': '1.0.0',
+            'langchain.task.name': task,
+        }
+
+        if len(args) > 0 and trace_input:
+            span_attributes['langchain.inputs'] = to_json_string(args)
+
+        attributes = FrameworkSpanAttributes(**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 trace_output:
+                    span.set_attribute(
+                        'langchain.outputs', to_json_string(result))
+
+                span.set_status(StatusCode.OK)
+                return result
+            except Exception as e:
+                # Record the exception in the span
+                span.record_exception(e)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    return traced_method
+
+
+def runnable_patch(method_name, task, tracer, version, trace_output=True, trace_input=True):
+    """
+    Wrapper function to trace a runnable
+    method_name: The name of the method to trace.
+    task: The name used to identify the type of task in `generic_patch`.
+    tracer: The tracer object used in `generic_patch`.
+    version: The version parameter used in `generic_patch`.
+    trace_output: Whether to trace the output of the patched methods.
+    trace_input: Whether to trace the input of the patched methods.
+    """
+    def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS['LANGCHAIN']
+        span_attributes = {
+            'langtrace.service.name': service_provider,
+            'langtrace.service.type': 'framework',
+            'langtrace.service.version': version,
+            'langtrace.version': '1.0.0',
+            'langchain.task.name': task,
+        }
+
+        if trace_input:
+            inputs = {}
+            args_list = []
+            if len(args) > 0:
+                for value in args:
+                    if isinstance(value, str):
+                        args_list.append(value)
+            inputs['args'] = args_list
+
+            for field, value in instance.steps.items() if hasattr(instance, "steps") and \
+                    isinstance(instance.steps, dict) else {}:
+                inputs[field] = value.__class__.__name__
+
+            span_attributes['langchain.inputs'] = to_json_string(inputs)
+
+        attributes = FrameworkSpanAttributes(**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 trace_output:
+                    outputs = {}
+                    if isinstance(result, dict):
+                        for field, value in result.items() if hasattr(result, "items") else {}:
+                            if isinstance(value, list):
+                                for item in value:
+                                    if item.__class__.__name__ == "Document":
+                                        outputs[field] = "Document"
+                                    else:
+                                        outputs[field] = item.__class__.__name__
+                            if isinstance(value, str):
+                                outputs[field] = value
+                    span.set_attribute(
+                        'langchain.outputs', to_json_string(outputs))
+                    if isinstance(result, str):
+                        span.set_attribute(
+                            'langchain.outputs', result)
+
+                span.set_status(StatusCode.OK)
+                return result
+            except Exception as e:
+                # Record the exception in the span
+                span.record_exception(e)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    return traced_method
+
+
+def clean_empty(d):
+    """Recursively remove empty lists, empty dicts, or None elements from a dictionary."""
+    if not isinstance(d, (dict, list)):
+        return d
+    if isinstance(d, list):
+        return [v for v in (clean_empty(v) for v in d) if v != [] and v is not None]
+    return {
+        k: v
+        for k, v in ((k, clean_empty(v)) for k, v in d.items())
+        if v is not None and v != {}
+    }
+
+
+def custom_serializer(obj):
+    """Fallback function to convert unserializable objects."""
+    if hasattr(obj, "__dict__"):
+        # Attempt to serialize custom objects by their __dict__ attribute.
+        return clean_empty(obj.__dict__)
+    else:
+        # For other types, just convert to string
+        return str(obj)
+
+
+def to_json_string(any_object):
+    """Converts any object to a JSON-parseable string, omitting empty or None values."""
+    cleaned_object = clean_empty(any_object)
+    return json.dumps(cleaned_object, default=custom_serializer, indent=2)

instrumentation/llamaindex/instrumentation.py

@@ -0,0 +1,73 @@
+import importlib.metadata
+from typing import Collection
+
+from langtrace.trace_attributes import LlamaIndexMethods
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from instrumentation.llamaindex.patch import generic_patch
+
+MODULES = [
+    "llama_index.core.query_pipeline.query",
+]
+
+
+class LlamaindexInstrumentation(BaseInstrumentor):
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["llama-index >= 0.10.0"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = importlib.metadata.version('llama-index')
+
+        wrap_function_wrapper(
+            'llama_index.core.base.base_query_engine',
+            'BaseQueryEngine.query',
+            generic_patch(LlamaIndexMethods.QUERYENGINE_QUERY.value,
+                          'query', tracer, version)
+        )
+        wrap_function_wrapper(
+            'llama_index.core.base.base_retriever',
+            'BaseRetriever.retrieve',
+            generic_patch(LlamaIndexMethods.RETRIEVER_RETRIEVE.value,
+                          'retrieve', tracer, version)
+        )
+        wrap_function_wrapper(
+            'llama_index.core.extractors.interface',
+            'BaseExtractor.extract',
+            generic_patch(
+                LlamaIndexMethods.BASEEXTRACTOR_EXTRACT.value, 'extract', tracer, version)
+        )
+        wrap_function_wrapper(
+            'llama_index.core.extractors.interface',
+            'BaseExtractor.aextract',
+            generic_patch(
+                LlamaIndexMethods.BASEEXTRACTOR_AEXTRACT.value, 'extract', tracer, version)
+        )
+        wrap_function_wrapper(
+            'llama_index.core.readers.file.base',
+            'SimpleDirectoryReader.load_data',
+            generic_patch(
+                LlamaIndexMethods.BASEREADER_LOADDATA.value, 'loaddata', tracer, version)
+        )
+        wrap_function_wrapper(
+            'llama_index.core.chat_engine.types',
+            'BaseChatEngine.chat',
+            generic_patch(
+                LlamaIndexMethods.CHATENGINE_CHAT.value, 'chat', tracer, version)
+        )
+        wrap_function_wrapper(
+            'llama_index.core.chat_engine.types',
+            'BaseChatEngine.achat',
+            generic_patch(
+                LlamaIndexMethods.CHATENGINE_ACHAT.value, 'chat', tracer, version)
+        )
+
+    def _instrument_module(self, module_name):
+        print(module_name)
+
+    def _uninstrument(self, **kwargs):
+        print(kwargs)

instrumentation/llamaindex/patch.py

@@ -0,0 +1,40 @@
+from langtrace.trace_attributes import FrameworkSpanAttributes
+from opentelemetry.trace import SpanKind, StatusCode
+from opentelemetry.trace.status import Status, StatusCode
+
+from instrumentation.constants import SERVICE_PROVIDERS
+
+
+def generic_patch(method, task, tracer, version):
+    def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS['LLAMAINDEX']
+        span_attributes = {
+            'langtrace.service.name': service_provider,
+            'langtrace.service.type': 'framework',
+            'langtrace.service.version': version,
+            'langtrace.version': '1.0.0',
+            'llamaindex.task.name': task,
+        }
+
+        attributes = FrameworkSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(method, 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)
+                span.set_status(StatusCode.OK)
+                return result
+            except Exception as e:
+                # Record the exception in the span
+                span.record_exception(e)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    return traced_method

instrumentation/openai/instrumentation.py

@@ -0,0 +1,41 @@
+import importlib.metadata
+from typing import Collection
+
+import openai
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from instrumentation.openai.patch import (chat_completions_create,
+                                          embeddings_create, images_generate)
+
+
+class OpenAIInstrumentation(BaseInstrumentor):
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["openai >= 0.27.0"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = importlib.metadata.version('openai')
+        wrap_function_wrapper(
+            'openai.resources.chat.completions',
+            'Completions.create',
+            chat_completions_create(
+                openai.chat.completions.create, version, tracer)
+        )
+        wrap_function_wrapper(
+            'openai.resources.images',
+            'Images.generate',
+            images_generate(openai.images.generate, version, tracer)
+        )
+        wrap_function_wrapper(
+            'openai.resources.embeddings',
+            'Embeddings.create',
+            embeddings_create(openai.embeddings.create,
+                              version, tracer)
+        )
+
+    def _uninstrument(self, **kwargs):
+        pass

instrumentation/openai/lib/apis.py

@@ -0,0 +1,16 @@
+from langtrace.trace_attributes import OpenAIMethods
+
+APIS = {
+    "CHAT_COMPLETION": {
+        "METHOD": OpenAIMethods.CHAT_COMPLETION.value,
+        "ENDPOINT": "/chat/completions",
+    },
+    "IMAGES_GENERATION": {
+        "METHOD": OpenAIMethods.IMAGES_GENERATION.value,
+        "ENDPOINT": "/images/generations",
+    },
+    "EMBEDDINGS_CREATE": {
+        "METHOD": OpenAIMethods.EMBEDDINGS_CREATE.value,
+        "ENDPOINT": "/embeddings",
+    },
+}

instrumentation/openai/lib/constants.py

@@ -0,0 +1,30 @@
+OPENAI_COST_TABLE = {
+    "gpt-4-0125-preview": {
+        "input": 0.01,
+        "output": 0.03,
+    },
+    "gpt-4-1106-preview": {
+        "input": 0.01,
+        "output": 0.03,
+    },
+    "gpt-4-1106-vision-preview": {
+        "input": 0.01,
+        "output": 0.03,
+    },
+    "gpt-4": {
+        "input": 0.03,
+        "output": 0.06,
+    },
+    "gpt-4-32k": {
+        "input": 0.06,
+        "output": 0.12,
+    },
+    "gpt-3.5-turbo-0125": {
+        "input": 0.0005,
+        "output": 0.0015,
+    },
+    "gpt-3.5-turbo-instruct": {
+        "input": 0.0015,
+        "output": 0.002,
+    },
+}

instrumentation/openai/patch.py

@@ -0,0 +1,209 @@
+import json
+
+from langtrace.trace_attributes import Event, LLMSpanAttributes
+from opentelemetry.trace import SpanKind, StatusCode
+from opentelemetry.trace.status import Status, StatusCode
+
+from instrumentation.constants import SERVICE_PROVIDERS
+from instrumentation.openai.lib.apis import APIS
+
+
+def images_generate(original_method, version, tracer):
+    def traced_method(wrapped, instance, args, kwargs):
+        base_url = str(instance._client._base_url) if hasattr(
+            instance, '_client') and hasattr(instance._client, '_base_url') else ""
+        service_provider = SERVICE_PROVIDERS['OPENAI']
+        span_attributes = {
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "llm",
+            "langtrace.service.version": version,
+            "langtrace.version": "1.0.0",
+            "url.full": base_url,
+            "llm.api": APIS["IMAGES_GENERATION"]["ENDPOINT"],
+            "llm.model": kwargs.get('model'),
+            "llm.stream": kwargs.get('stream'),
+            "llm.prompts": json.dumps([kwargs.get('prompt', [])])
+        }
+
+        attributes = LLMSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(APIS["IMAGES_GENERATION"]["METHOD"], 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 = original_method(*args, **kwargs)
+                if kwargs.get('stream') is False or kwargs.get('stream') is None:
+                    data = result.data[0] if hasattr(
+                        result, 'data') and len(result.data) > 0 else {}
+                    response = [{
+                        "url": data.url if hasattr(data, 'url') else "",
+                        "revised_prompt": data.revised_prompt if hasattr(data, 'revised_prompt') else "",
+                    }]
+                    span.set_attribute(
+                        "llm.responses", json.dumps(response))
+
+                span.set_status(StatusCode.OK)
+                return result
+            except Exception as e:
+                # Record the exception in the span
+                span.record_exception(e)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    return traced_method
+
+
+def chat_completions_create(original_method, version, tracer):
+    def traced_method(wrapped, instance, args, kwargs):
+        base_url = str(instance._client._base_url) if hasattr(
+            instance, '_client') and hasattr(instance._client, '_base_url') else ""
+        service_provider = SERVICE_PROVIDERS['OPENAI']
+        span_attributes = {
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "llm",
+            "langtrace.service.version": version,
+            "langtrace.version": "1.0.0",
+            "url.full": base_url,
+            "llm.api": APIS["CHAT_COMPLETION"]["ENDPOINT"],
+            "llm.model": kwargs.get('model'),
+            "llm.prompts": json.dumps(kwargs.get('messages', [])),
+            "llm.stream": kwargs.get('stream'),
+        }
+
+        attributes = LLMSpanAttributes(**span_attributes)
+
+        if kwargs.get('temperature') is not None:
+            attributes.llm_temperature = kwargs.get('temperature')
+        if kwargs.get('top_p') is not None:
+            attributes.llm_top_p = kwargs.get('top_p')
+        if kwargs.get('user') is not None:
+            attributes.llm_user = kwargs.get('user')
+
+        with tracer.start_as_current_span(APIS["CHAT_COMPLETION"]["METHOD"], 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 = original_method(*args, **kwargs)
+                if kwargs.get('stream') is False:
+                    if hasattr(result, 'choices') and result.choices is not None:
+                        responses = [
+                            {
+                                "message": choice.message.content if choice.message and choice.message.content else "",
+                                **({"content_filter_results": choice["content_filter_results"]} if "content_filter_results" in choice else {})
+                            }
+                            for choice in result.choices
+                        ]
+                    else:
+                        responses = []
+                    span.set_attribute("llm.responses", json.dumps(responses))
+
+                    if hasattr(result, 'system_fingerprint') and result.system_fingerprint is not None:
+                        span.set_attribute(
+                            "llm.system.fingerprint", result.system_fingerprint)
+
+                    # Get the usage
+                    if hasattr(result, 'usage') and result.usage is not None:
+                        usage = result.usage
+                        if usage is not None:
+                            usage_dict = {
+                                "prompt_tokens": result.usage.prompt_tokens,
+                                "completion_tokens": usage.completion_tokens,
+                                "total_tokens": usage.total_tokens
+                            }
+                            span.set_attribute(
+                                "llm.token.counts", json.dumps(usage_dict))
+
+                    span.set_status(StatusCode.OK)
+                    return result
+                else:
+                    result_content = []
+                    span.add_event(Event.STREAM_START.value)
+
+                    for chunk in result:
+                        # Assuming `chunk` has a structure similar to what OpenAI might return,
+                        # adjust the access accordingly based on actual response structure.
+                        if hasattr(chunk, 'choices') and chunk.choices is not None:
+                            content = [
+                                choice.delta.content if choice.delta and choice.delta.content else ""
+                                for choice in chunk.choices
+                            ]
+                        else:
+                            content = []
+                        span.add_event(Event.STREAM_OUTPUT.value, {
+                            "response": "".join(content)
+                        })
+                        result_content.append(
+                            content[0] if len(content) > 0 else "")
+                    span.add_event(Event.STREAM_END.value)
+                    span.set_attribute("llm.responses", json.dumps(
+                        {"message": {"role": "assistant", "content": "".join(result_content)}}))
+
+            except Exception as e:
+                # Record the exception in the span
+                span.record_exception(e)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    # return the wrapped method
+    return traced_method
+
+
+def embeddings_create(original_method, version, tracer):
+    def traced_method(wrapped, instance, args, kwargs):
+        base_url = str(instance._client._base_url) if hasattr(
+            instance, '_client') and hasattr(instance._client, '_base_url') else ""
+
+        service_provider = SERVICE_PROVIDERS['OPENAI']
+        span_attributes = {
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "llm",
+            "langtrace.service.version": version,
+            "langtrace.version": "1.0.0",
+            "url.full": base_url,
+            "llm.api": APIS["EMBEDDINGS_CREATE"]["ENDPOINT"],
+            "llm.model": kwargs.get('model'),
+            "llm.prompts": "",
+        }
+
+        attributes = LLMSpanAttributes(**span_attributes)
+        kwargs.get('encoding_format')
+
+        if kwargs.get('encoding_format') is not None:
+            attributes.llm_encoding_format = kwargs.get('encoding_format')
+        if kwargs.get('dimensions') is not None:
+            attributes["llm.dimensions"] = kwargs.get('dimensions')
+        if kwargs.get('user') is not None:
+            attributes["llm.user"] = kwargs.get('user')
+
+        with tracer.start_as_current_span(APIS["EMBEDDINGS_CREATE"]["METHOD"], 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 = original_method(*args, **kwargs)
+                span.set_status(StatusCode.OK)
+                return result
+            except Exception as e:
+                # Record the exception in the span
+                span.record_exception(e)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    return traced_method

instrumentation/pinecone/instrumentation.py

@@ -0,0 +1,43 @@
+import importlib.metadata
+from typing import Collection
+
+import pinecone
+from langtrace.trace_attributes import PineconeMethods
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from instrumentation.pinecone.lib.apis import APIS
+from instrumentation.pinecone.patch import generic_patch
+
+
+class PineconeInstrumentation(BaseInstrumentor):
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["pinecone-client >= 3.1.0"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = importlib.metadata.version('pinecone-client')
+        for operation_name, details in APIS.items():
+            method_ref = details["METHOD"]
+            method = None
+            if method_ref is PineconeMethods.UPSERT.value:
+                method = pinecone.Index.upsert
+            elif method_ref is PineconeMethods.QUERY.value:
+                method = pinecone.Index.query
+            elif method_ref is PineconeMethods.DELETE.value:
+                method = pinecone.Index.delete
+            operation = details["OPERATION"]
+
+            # Dynamically creating the patching call
+            wrap_function_wrapper(
+                'pinecone.data.index',
+                f'Index.{operation}',
+                generic_patch(method, operation_name,
+                              version, tracer)
+            )
+
+    def _uninstrument(self, **kwargs):
+        print(kwargs)

instrumentation/pinecone/lib/apis.py

@@ -0,0 +1,19 @@
+from langtrace.trace_attributes import PineconeMethods
+
+APIS = {
+    "UPSERT": {
+        "METHOD": PineconeMethods.UPSERT.value,
+        "ENDPOINT": "/vectors/upsert",
+        "OPERATION": "upsert",
+    },
+    "QUERY": {
+        "METHOD": PineconeMethods.QUERY.value,
+        "ENDPOINT": "/query",
+        "OPERATION": "query",
+    },
+    "DELETE": {
+        "METHOD": PineconeMethods.DELETE.value,
+        "ENDPOINT": "/vectors/delete",
+        "OPERATION": "delete",
+    }
+}