langtrace-python-sdk 2.1.16__py3-none-any.whl → 2.1.19__py3-none-any.whl

langtrace_python_sdk/instrumentation/ollama/instrumentation.py

@@ -0,0 +1,58 @@
+"""
+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 wrapt import wrap_function_wrapper as _W
+from typing import Collection
+from importlib_metadata import version as v
+from langtrace_python_sdk.constants.instrumentation.ollama import APIS
+from .patch import generic_patch, ageneric_patch
+
+
+class OllamaInstrumentor(BaseInstrumentor):
+    """
+    The OllamaInstrumentor class represents the Ollama instrumentation"""
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["ollama >= 0.2.0, < 1"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = v("ollama")
+        for operation_name, details in APIS.items():
+            operation = details["METHOD"]
+            # Dynamically creating the patching call
+            _W(
+                "ollama._client",
+                f"Client.{operation}",
+                generic_patch(operation_name, version, tracer),
+            )
+
+            _W(
+                "ollama._client",
+                f"AsyncClient.{operation}",
+                ageneric_patch(operation_name, version, tracer),
+            )
+            _W(
+                "ollama",
+                f"{operation}",
+                generic_patch(operation_name, version, tracer),
+            )
+
+    def _uninstrument(self, **kwargs):
+        pass

langtrace_python_sdk/instrumentation/ollama/patch.py

@@ -0,0 +1,215 @@
+from langtrace_python_sdk.constants.instrumentation.ollama import APIS
+from importlib_metadata import version as v
+from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
+from langtrace_python_sdk.utils import set_span_attribute
+from langtrace_python_sdk.utils.silently_fail import silently_fail
+from langtrace_python_sdk.constants.instrumentation.common import (
+    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
+    SERVICE_PROVIDERS,
+)
+from opentelemetry import baggage
+from langtrace.trace_attributes import LLMSpanAttributes, Event
+from opentelemetry.trace import SpanKind
+import json
+from opentelemetry.trace.status import Status, StatusCode
+
+
+def generic_patch(operation_name, 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 ""
+        )
+        api = APIS[operation_name]
+        service_provider = SERVICE_PROVIDERS["OLLAMA"]
+        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
+        span_attributes = {
+            "langtrace.sdk.name": "langtrace-python-sdk",
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "llm",
+            "langtrace.service.version": version,
+            "langtrace.version": v(LANGTRACE_SDK_NAME),
+            "llm.model": kwargs.get("model"),
+            "llm.stream": kwargs.get("stream"),
+            "url.full": base_url,
+            "llm.api": api["METHOD"],
+            **(extra_attributes if extra_attributes is not None else {}),
+        }
+
+        attributes = LLMSpanAttributes(**span_attributes)
+        with tracer.start_as_current_span(
+            f'ollama.{api["METHOD"]}', kind=SpanKind.CLIENT
+        ) as span:
+            _set_input_attributes(span, kwargs, attributes)
+
+            try:
+                result = wrapped(*args, **kwargs)
+                if result:
+                    if span.is_recording():
+
+                        if kwargs.get("stream"):
+                            return _handle_streaming_response(
+                                span, result, api["METHOD"]
+                            )
+
+                        _set_response_attributes(span, result)
+                        span.set_status(Status(StatusCode.OK))
+
+                span.end()
+                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 ageneric_patch(operation_name, version, tracer):
+    async def traced_method(wrapped, instance, args, kwargs):
+        api = APIS[operation_name]
+        service_provider = SERVICE_PROVIDERS["OLLAMA"]
+        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
+        span_attributes = {
+            "langtrace.sdk.name": "langtrace-python-sdk",
+            "langtrace.service.name": service_provider,
+            "url.full": "",
+            "llm.api": "",
+            "langtrace.service.type": "llm",
+            "langtrace.service.version": version,
+            "langtrace.version": v(LANGTRACE_SDK_NAME),
+            "llm.model": kwargs.get("model"),
+            "llm.stream": kwargs.get("stream"),
+            **(extra_attributes if extra_attributes is not None else {}),
+        }
+
+        attributes = LLMSpanAttributes(**span_attributes)
+        with tracer.start_as_current_span(api["METHOD"], kind=SpanKind.CLIENT) as span:
+            _set_input_attributes(span, kwargs, attributes)
+            try:
+                result = await wrapped(*args, **kwargs)
+                if result:
+                    if span.is_recording():
+                        if kwargs.get("stream"):
+                            return _ahandle_streaming_response(
+                                span, result, api["METHOD"]
+                            )
+
+                        _set_response_attributes(span, result)
+                        span.set_status(Status(StatusCode.OK))
+                span.end()
+                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
+
+
+@silently_fail
+def _set_response_attributes(span, response):
+
+    input_tokens = response.get("prompt_eval_count") or 0
+    output_tokens = response.get("eval_count") or 0
+    total_tokens = input_tokens + output_tokens
+    usage_dict = {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "total_tokens": total_tokens,
+    }
+
+    if total_tokens > 0:
+        set_span_attribute(span, "llm.token.counts", json.dumps(usage_dict))
+    set_span_attribute(span, "llm.finish_reason", response.get("done_reason"))
+
+    if "message" in response:
+        set_span_attribute(span, "llm.responses", json.dumps([response.get("message")]))
+
+    if "response" in response:
+        set_span_attribute(
+            span, "llm.responses", json.dumps([response.get("response")])
+        )
+
+
+@silently_fail
+def _set_input_attributes(span, kwargs, attributes):
+    for field, value in attributes.model_dump(by_alias=True).items():
+        set_span_attribute(span, field, value)
+
+    if "messages" in kwargs:
+        set_span_attribute(
+            span,
+            "llm.prompts",
+            json.dumps([kwargs.get("messages", [])]),
+        )
+    if "prompt" in kwargs:
+        set_span_attribute(
+            span,
+            "llm.prompts",
+            json.dumps([{"role": "user", "content": kwargs.get("prompt", [])}]),
+        )
+
+
+def _handle_streaming_response(span, response, api):
+    accumulated_tokens = None
+    if api == "chat":
+        accumulated_tokens = {"message": {"content": "", "role": ""}}
+    if api == "completion":
+        accumulated_tokens = {"response": ""}
+
+    span.add_event(Event.STREAM_START.value)
+    try:
+        for chunk in response:
+            if api == "chat":
+                accumulated_tokens["message"]["content"] += chunk["message"]["content"]
+                accumulated_tokens["message"]["role"] = chunk["message"]["role"]
+            if api == "generate":
+                accumulated_tokens["response"] += chunk["response"]
+
+        _set_response_attributes(span, chunk | accumulated_tokens)
+    finally:
+        # Finalize span after processing all chunks
+        span.add_event(Event.STREAM_END.value)
+        span.set_status(StatusCode.OK)
+        span.end()
+
+    return response
+
+
+async def _ahandle_streaming_response(span, response, api):
+    accumulated_tokens = None
+    if api == "chat":
+        accumulated_tokens = {"message": {"content": "", "role": ""}}
+    if api == "completion":
+        accumulated_tokens = {"response": ""}
+
+    span.add_event(Event.STREAM_START.value)
+    try:
+        async for chunk in response:
+            if api == "chat":
+                accumulated_tokens["message"]["content"] += chunk["message"]["content"]
+                accumulated_tokens["message"]["role"] = chunk["message"]["role"]
+            if api == "generate":
+                accumulated_tokens["response"] += chunk["response"]
+
+        _set_response_attributes(span, chunk | accumulated_tokens)
+    finally:
+        # Finalize span after processing all chunks
+        span.add_event(Event.STREAM_END.value)
+        span.set_status(StatusCode.OK)
+        span.end()
+
+    return response

langtrace_python_sdk/instrumentation/openai/patch.py

@@ -18,10 +18,10 @@ import json
 
 from importlib_metadata import version as v
 from langtrace.trace_attributes import Event, LLMSpanAttributes
-from opentelemetry import baggage
+from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
-
+from opentelemetry.trace.propagation import set_span_in_context
 from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
@@ -64,7 +64,9 @@ def images_generate(original_method, version, tracer):
         attributes = LLMSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            APIS["IMAGES_GENERATION"]["METHOD"], kind=SpanKind.CLIENT
+            APIS["IMAGES_GENERATION"]["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
         ) as span:
             for field, value in attributes.model_dump(by_alias=True).items():
                 if value is not None:
@@ -141,7 +143,9 @@ def async_images_generate(original_method, version, tracer):
         attributes = LLMSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            APIS["IMAGES_GENERATION"]["METHOD"], kind=SpanKind.CLIENT
+            APIS["IMAGES_GENERATION"]["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
         ) as span:
             items = attributes.model_dump(by_alias=True).items()
             for field, value in items:
@@ -226,7 +230,9 @@ def images_edit(original_method, version, tracer):
         attributes = LLMSpanAttributes(**span_attributes)
 
         with tracer.start_as_current_span(
-            APIS["IMAGES_EDIT"]["METHOD"], kind=SpanKind.CLIENT
+            APIS["IMAGES_EDIT"]["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
         ) as span:
             for field, value in attributes.model_dump(by_alias=True).items():
                 if value is not None:
@@ -349,7 +355,9 @@ def chat_completions_create(original_method, version, tracer):
         # with tracer.start_as_current_span(APIS["CHAT_COMPLETION"]["METHOD"],
         #                                   kind=SpanKind.CLIENT) as span:
         span = tracer.start_span(
-            APIS["CHAT_COMPLETION"]["METHOD"], kind=SpanKind.CLIENT
+            APIS["CHAT_COMPLETION"]["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
         )
         for field, value in attributes.model_dump(by_alias=True).items():
             if value is not None:
@@ -833,7 +841,9 @@ def embeddings_create(original_method, version, tracer):
             attributes["llm.user"] = kwargs.get("user")
 
         with tracer.start_as_current_span(
-            APIS["EMBEDDINGS_CREATE"]["METHOD"], kind=SpanKind.CLIENT
+            APIS["EMBEDDINGS_CREATE"]["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
         ) as span:
 
             for field, value in attributes.model_dump(by_alias=True).items():
@@ -898,7 +908,9 @@ def async_embeddings_create(original_method, version, tracer):
             attributes["llm.user"] = kwargs.get("user")
 
         with tracer.start_as_current_span(
-            APIS["EMBEDDINGS_CREATE"]["METHOD"], kind=SpanKind.CLIENT
+            APIS["EMBEDDINGS_CREATE"]["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
         ) as span:
 
             async for field, value in attributes.model_dump(by_alias=True).items():

langtrace_python_sdk/instrumentation/pinecone/instrumentation.py

@@ -18,7 +18,6 @@ import importlib.metadata
 import logging
 from typing import Collection
 
-from langtrace.trace_attributes import PineconeMethods
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.trace import get_tracer
 from wrapt import wrap_function_wrapper

langtrace_python_sdk/instrumentation/pinecone/patch.py

@@ -17,10 +17,10 @@ limitations under the License.
 import json
 
 from langtrace.trace_attributes import DatabaseSpanAttributes
-from opentelemetry import baggage
+from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
-
+from opentelemetry.trace.propagation import set_span_in_context
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
     SERVICE_PROVIDERS,
@@ -56,7 +56,11 @@ def generic_patch(operation_name, version, tracer):
 
         attributes = DatabaseSpanAttributes(**span_attributes)
 
-        with tracer.start_as_current_span(api["METHOD"], kind=SpanKind.CLIENT) as span:
+        with tracer.start_as_current_span(
+            api["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
 
             if span.is_recording():
                 set_span_attribute(span, "server.address", instance._config.host)

langtrace_python_sdk/instrumentation/qdrant/patch.py

@@ -18,9 +18,10 @@ 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
-from opentelemetry import baggage
+from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
+from opentelemetry.trace.propagation import set_span_in_context
 
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
@@ -56,7 +57,11 @@ def collection_patch(method, version, tracer):
 
         attributes = DatabaseSpanAttributes(**span_attributes)
 
-        with tracer.start_as_current_span(api["METHOD"], kind=SpanKind.CLIENT) as span:
+        with tracer.start_as_current_span(
+            api["METHOD"],
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
             collection_name = kwargs.get("collection_name") or args[0]
             operation = api["OPERATION"]
             set_span_attributes(span, "db.collection.name", collection_name)

langtrace_python_sdk/instrumentation/weaviate/patch.py

@@ -17,10 +17,10 @@ limitations under the License.
 import json
 
 from langtrace.trace_attributes import DatabaseSpanAttributes
-from opentelemetry import baggage
+from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
-
+from opentelemetry.trace.propagation import set_span_in_context
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
     SERVICE_PROVIDERS,
@@ -131,7 +131,11 @@ def create_traced_method(method_name, version, tracer, get_collection_name=None)
 
         attributes = DatabaseSpanAttributes(**span_attributes)
 
-        with tracer.start_as_current_span(method_name, kind=SpanKind.CLIENT) as span:
+        with tracer.start_as_current_span(
+            method_name,
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
             for field, value in attributes.model_dump(by_alias=True).items():
                 if value is not None:
                     span.set_attribute(field, value)

langtrace_python_sdk/langtrace.py

@@ -19,7 +19,12 @@ from typing import Optional
 from langtrace_python_sdk.constants.exporter.langtrace_exporter import (
     LANGTRACE_REMOTE_URL,
 )
-from langtrace_python_sdk.types import DisableInstrumentations, InstrumentationType
+from langtrace_python_sdk.types import (
+    DisableInstrumentations,
+    InstrumentationType,
+    InstrumentationMethods,
+)
+from langtrace_python_sdk.utils.langtrace_sampler import LangtraceSampler
 from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import (
@@ -45,6 +50,7 @@ from langtrace_python_sdk.instrumentation import (
     PineconeInstrumentation,
     QdrantInstrumentation,
     WeaviateInstrumentation,
+    OllamaInstrumentor,
 )
 from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
 from colorama import Fore
@@ -59,6 +65,7 @@ def init(
     custom_remote_exporter=None,
     api_host: Optional[str] = LANGTRACE_REMOTE_URL,
     disable_instrumentations: Optional[DisableInstrumentations] = None,
+    disable_tracing_for_functions: Optional[InstrumentationMethods] = None,
 ):
 
     host = (
@@ -66,7 +73,11 @@ def init(
     )
     check_if_sdk_is_outdated()
     print(Fore.GREEN + "Initializing Langtrace SDK.." + Fore.RESET)
-    provider = TracerProvider(resource=Resource.create({"service.name": sys.argv[0]}))
+    sampler = LangtraceSampler(disabled_methods=disable_tracing_for_functions)
+    provider = TracerProvider(
+        resource=Resource.create({"service.name": sys.argv[0]}),
+        sampler=sampler,
+    )
 
     remote_write_exporter = (
         LangTraceExporter(api_key=api_key, api_host=host)
@@ -96,6 +107,7 @@ def init(
         "cohere": CohereInstrumentation(),
         "weaviate": WeaviateInstrumentation(),
         "sqlalchemy": SQLAlchemyInstrumentor(),
+        "ollama": OllamaInstrumentor(),
     }
 
     init_instrumentations(disable_instrumentations, all_instrumentations)

langtrace_python_sdk/types/__init__.py

@@ -1,4 +1,4 @@
-from typing import List, TypedDict
+from typing import List, Literal, TypedDict
 from enum import Enum
 
 
@@ -16,6 +16,7 @@ class InstrumentationType(Enum):
     LANGCHAIN_COMMUNITY = "langchain_community"
     LANGGRAPH = "langgraph"
     WEAVIATE = "weaviate"
+    OLLAMA = "ollama"
 
     @staticmethod
     def from_string(value: str):
@@ -28,3 +29,78 @@ class InstrumentationType(Enum):
 class DisableInstrumentations(TypedDict, total=False):
     all_except: List[InstrumentationType]
     only: List[InstrumentationType]
+
+
+class VendorMethods(TypedDict):
+    PineconeMethods = Literal[
+        "pinecone.index.upsert", "pinecone.index.query", "pinecone.index.delete"
+    ]
+    AnthropicMethods = Literal["anthropic.messages.create"]
+    GroqMethods = Literal["groq.chat.completions.create"]
+    OpenaiMethods = Literal[
+        "openai.chat.completions.create",
+        "openai.embeddings.create",
+        "openai.images.generate",
+        "openai.images.edit",
+    ]
+
+    ChromadbMethods = Literal[
+        "chromadb.collection.add",
+        "chromadb.collection.query",
+        "chromadb.collection.delete",
+        "chromadb.collection.peek",
+        "chromadb.collection.update",
+        "chromadb.collection.upsert",
+        "chromadb.collection.modify",
+        "chromadb.collection.count",
+    ]
+    QdrantMethods = Literal[
+        "qdrantdb.add",
+        "qdrantdb.get_collection",
+        "qdrantdb.get_collections",
+        "qdrantdb.query",
+        "qdrantdb.query_batch",
+        "qdrantdb.delete",
+        "qdrantdb.discover",
+        "qdrantdb.discover_batch",
+        "qdrantdb.recommend",
+        "qdrantdb.recommend_batch",
+        "qdrantdb.retrieve",
+        "qdrantdb.search",
+        "qdrantdb.search_batch",
+        "qdrantdb.upsert",
+        "qdrantdb.count",
+        "qdrantdb.update_collection",
+        "qdrantdb.update_vectors",
+    ]
+
+    CohereMethods = Literal[
+        "cohere.client.chat",
+        "cohere.client.embed",
+        "cohere.client.rerank",
+        "cohere.client.chat_stream",
+    ]
+
+    LlamaIndexMethods = Literal[
+        "llamaindex.OpenAI.chat",
+        "llamaindex.RetrieverQueryEngine.query",
+        "llamaindex.VectorIndexRetriever.retrieve",
+        "llamaindex.SimpleVectorStore.query",
+        "llamaindex.RetrieverQueryEngine.retrieve",
+    ]
+
+
+class InstrumentationMethods(TypedDict):
+    open_ai: List[VendorMethods.OpenaiMethods]
+    groq: List[VendorMethods.GroqMethods]
+    pinecone: List[VendorMethods.PineconeMethods]
+    llamaindex: List[VendorMethods.LlamaIndexMethods]
+    chromadb: List[VendorMethods.ChromadbMethods]
+    qdrant: List[VendorMethods.QdrantMethods]
+    langchain: List[str]
+    langchain_core: List[str]
+    langchain_community: List[str]
+    langgraph: List[str]
+    anthropic: List[VendorMethods.AnthropicMethods]
+    cohere: List[VendorMethods.CohereMethods]
+    weaviate: List[str]

langtrace_python_sdk/utils/langtrace_sampler.py

@@ -0,0 +1,56 @@
+from typing import Optional, Sequence
+from opentelemetry.sdk.trace.sampling import (
+    Sampler,
+    Decision,
+    SamplingResult,
+)
+from opentelemetry.trace import SpanKind, Link, TraceState, TraceFlags
+from opentelemetry.util.types import Attributes
+from opentelemetry.context import Context
+from opentelemetry import trace
+
+
+class LangtraceSampler(Sampler):
+    _disabled_methods_names: set
+
+    def __init__(
+        self,
+        disabled_methods: dict,
+    ):
+        self._disabled_methods_names = set()
+        if disabled_methods:
+            for _, methods in disabled_methods.items():
+                for method in methods:
+                    self._disabled_methods_names.add(method)
+
+    def should_sample(
+        self,
+        parent_context: Optional[Context],
+        trace_id: int,
+        name: str,
+        kind: Optional[SpanKind] = None,
+        attributes: Attributes = None,
+        links: Optional[Sequence[Link]] = None,
+        trace_state: Optional["TraceState"] = None,
+    ) -> SamplingResult:
+
+        parent_span = trace.get_current_span(parent_context)
+        parent_span_context = parent_span.get_span_context()
+
+        if not self._disabled_methods_names:
+            return SamplingResult(decision=Decision.RECORD_AND_SAMPLE)
+
+        if parent_context:
+            if (
+                parent_span_context.span_id != 0
+                and parent_span_context.trace_flags != TraceFlags.SAMPLED
+            ):
+                return SamplingResult(decision=Decision.DROP)
+
+        if name in self._disabled_methods_names:
+            return SamplingResult(decision=Decision.DROP)
+
+        return SamplingResult(decision=Decision.RECORD_AND_SAMPLE)
+
+    def get_description(self):
+        return "Langtrace Sampler"

langtrace_python_sdk/utils/with_root_span.py

@@ -23,6 +23,7 @@ from typing import Optional
 import requests
 from opentelemetry import baggage, context, trace
 from opentelemetry.trace import SpanKind
+from opentelemetry.trace.propagation import set_span_in_context
 
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
@@ -48,7 +49,11 @@ def with_langtrace_root_span(
             span_id = None
             trace_id = None
 
-            with tracer.start_as_current_span(operation_name, kind=kind) as span:
+            with tracer.start_as_current_span(
+                operation_name,
+                kind=kind,
+                context=set_span_in_context(trace.get_current_span()),
+            ) as span:
                 span_id = str(span.get_span_context().span_id)
                 trace_id = str(span.get_span_context().trace_id)
 
@@ -66,7 +71,11 @@ def with_langtrace_root_span(
             trace_id = None
             tracer = trace.get_tracer(__name__)
             operation_name = name if name else func.__name__
-            with tracer.start_as_current_span(operation_name, kind=kind) as span:
+            with tracer.start_as_current_span(
+                operation_name,
+                kind=kind,
+                context=set_span_in_context(trace.get_current_span()),
+            ) as span:
                 span_id = span.get_span_context().span_id
                 trace_id = span.get_span_context().trace_id
                 if (

langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "2.1.16"
+__version__ = "2.1.19"