langtrace-python-sdk 2.1.28__py3-none-any.whl → 2.2.1__py3-none-any.whl

langtrace_python_sdk/instrumentation/qdrant/patch.py

@@ -17,7 +17,7 @@ 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
+from langtrace_python_sdk.utils import set_span_attribute
 from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
@@ -64,7 +64,7 @@ def collection_patch(method, version, tracer):
         ) as span:
             collection_name = kwargs.get("collection_name") or args[0]
             operation = api["OPERATION"]
-            set_span_attributes(span, "db.collection.name", collection_name)
+            set_span_attribute(span, "db.collection.name", collection_name)
 
             if operation == "add":
                 _set_upload_attributes(span, args, kwargs, "documents")
@@ -111,7 +111,7 @@ def _set_upsert_attributes(span, args, kwargs):
     else:
         # In case of using Batch.
         length = len(points.ids)
-    set_span_attributes(span, "db.upsert.points_count", length)
+    set_span_attribute(span, "db.upsert.points_count", length)
 
 
 @silently_fail
@@ -123,16 +123,16 @@ def _set_upload_attributes(span, args, kwargs, field):
         # In case of using Batch.
         length = len(docs.ids)
 
-    set_span_attributes(span, f"db.upload.{field}_count", length)
+    set_span_attribute(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)
+    set_span_attribute(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))
+    set_span_attribute(span, f"db.{method}.requests_count", len(requests))

langtrace_python_sdk/instrumentation/vertexai/__init__.py

@@ -0,0 +1,3 @@
+from .instrumentation import VertexAIInstrumentation
+
+__all__ = ["VertexAIInstrumentation"]

langtrace_python_sdk/instrumentation/vertexai/instrumentation.py

@@ -0,0 +1,33 @@
+from typing import Collection
+from importlib_metadata import version as v
+from langtrace_python_sdk.constants.instrumentation.vertexai import APIS
+from wrapt import wrap_function_wrapper as _W
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from .patch import patch_vertexai
+
+
+class VertexAIInstrumentation(BaseInstrumentor):
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["google-cloud-aiplatform >= 1.0.0"]
+
+    def _instrument(self, **kwargs):
+        trace_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", trace_provider)
+        version = v("google-cloud-aiplatform")
+
+        for _, api_config in APIS.items():
+
+            module = api_config.get("module")
+            operation = api_config.get("operation")
+            method = api_config.get("method")
+            name = f"{method}.{operation}"
+
+            _W(
+                module=module,
+                name=name,
+                wrapper=patch_vertexai(name, version, tracer),
+            )
+
+    def _uninstrument(self, **kwargs):
+        pass

langtrace_python_sdk/instrumentation/vertexai/patch.py

@@ -0,0 +1,131 @@
+import types
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
+
+
+from langtrace_python_sdk.utils.llm import (
+    calculate_prompt_tokens,
+    get_extra_attributes,
+    get_langtrace_attributes,
+    get_llm_request_attributes,
+    get_llm_url,
+    set_event_completion,
+    set_span_attributes,
+    set_usage_attributes,
+    StreamWrapper,
+)
+from langtrace.trace_attributes import LLMSpanAttributes, SpanAttributes
+from langtrace_python_sdk.utils.silently_fail import silently_fail
+from opentelemetry.trace import Tracer, SpanKind, Span
+from opentelemetry import trace
+from opentelemetry.trace.propagation import set_span_in_context
+from opentelemetry.trace.status import Status, StatusCode
+import json
+
+
+def patch_vertexai(name, version, tracer: Tracer):
+    def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS["VERTEXAI"]
+        prompts = serialize_prompts(args, kwargs)
+        span_attributes = {
+            **get_langtrace_attributes(version, service_provider),
+            **get_llm_request_attributes(
+                kwargs,
+                prompts=prompts,
+                model=get_llm_model(instance),
+            ),
+            **get_llm_url(instance),
+            SpanAttributes.LLM_PATH: "",
+            **get_extra_attributes(),
+        }
+        attributes = LLMSpanAttributes(**span_attributes)
+        span = tracer.start_span(
+            name=name,
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        )
+
+        try:
+            set_span_attributes(span, attributes)
+            result = wrapped(*args, **kwargs)
+            if is_streaming_response(result):
+                prompt_tokens = 0
+                for message in kwargs.get("message", {}):
+                    prompt_tokens += calculate_prompt_tokens(
+                        json.dumps(message), kwargs.get("model")
+                    )
+                return StreamWrapper(
+                    stream=result, span=span, prompt_tokens=prompt_tokens
+                )
+            else:
+                set_response_attributes(span, result)
+            span.set_status(StatusCode.OK)
+            span.end()
+            return result
+        except Exception as error:
+            span.record_exception(error)
+            span.set_status(Status(StatusCode.ERROR, str(error)))
+            span.end()
+            raise
+
+    return traced_method
+
+
+@silently_fail
+def set_response_attributes(span: Span, result):
+
+    if hasattr(result, "text"):
+        set_event_completion(span, [{"role": "assistant", "content": result.text}])
+
+    if hasattr(result, "usage_metadata"):
+        usage = result.usage_metadata
+        input_tokens = usage.prompt_token_count
+        output_tokens = usage.candidates_token_count
+
+        set_usage_attributes(
+            span, {"input_tokens": input_tokens, "output_tokens": output_tokens}
+        )
+
+    if hasattr(result, "_prediction_response"):
+        usage = result._prediction_response.metadata.get("tokenMetadata")
+        input_tokens = usage.get("inputTokenCount").get("totalTokens")
+        output_tokens = usage.get("outputTokenCount").get("totalTokens")
+        set_usage_attributes(
+            span, {"input_tokens": input_tokens, "output_tokens": output_tokens}
+        )
+
+
+def is_streaming_response(response):
+    return isinstance(response, types.GeneratorType) or isinstance(
+        response, types.AsyncGeneratorType
+    )
+
+
+def get_llm_model(instance):
+    llm_model = "unknown"
+    if hasattr(instance, "_model_id"):
+        llm_model = instance._model_id
+    if hasattr(instance, "_model_name"):
+        llm_model = instance._model_name.replace("publishers/google/models/", "")
+    return llm_model
+
+
+def serialize_prompts(args, kwargs):
+    prompt = ""
+    if args is not None and len(args) > 0:
+        for arg in args:
+            if isinstance(arg, str):
+                prompt = f"{prompt}{arg}\n"
+            elif isinstance(arg, list):
+                for subarg in arg:
+                    if type(subarg).__name__ == "Part":
+                        prompt = f"{prompt}{json.dumps(subarg.to_dict())}\n"
+                    else:
+                        prompt = f"{prompt}{subarg}\n"
+    else:
+        prompt = [
+            {
+                "role": "user",
+                "content": kwargs.get("prompt") or kwargs.get("message"),
+            }
+        ]
+    return prompt

langtrace_python_sdk/langtrace.py

@@ -53,6 +53,8 @@ from langtrace_python_sdk.instrumentation import (
     WeaviateInstrumentation,
     OllamaInstrumentor,
     DspyInstrumentation,
+    VertexAIInstrumentation,
+    GeminiInstrumentation,
 )
 from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
 from colorama import Fore
@@ -68,6 +70,7 @@ def init(
     api_host: Optional[str] = LANGTRACE_REMOTE_URL,
     disable_instrumentations: Optional[DisableInstrumentations] = None,
     disable_tracing_for_functions: Optional[InstrumentationMethods] = None,
+    service_name: Optional[str] = None,
 ):
 
     host = (
@@ -77,7 +80,7 @@ def init(
     print(Fore.GREEN + "Initializing Langtrace SDK.." + Fore.RESET)
     sampler = LangtraceSampler(disabled_methods=disable_tracing_for_functions)
     provider = TracerProvider(
-        resource=Resource.create({"service.name": sys.argv[0]}),
+        resource=Resource.create({"service.name": service_name or sys.argv[0]}),
         sampler=sampler,
     )
 
@@ -112,6 +115,8 @@ def init(
         "ollama": OllamaInstrumentor(),
         "dspy": DspyInstrumentation(),
         "crewai": CrewAIInstrumentation(),
+        "vertexai": VertexAIInstrumentation(),
+        "gemini": GeminiInstrumentation(),
     }
 
     init_instrumentations(disable_instrumentations, all_instrumentations)
@@ -142,6 +147,7 @@ def init_instrumentations(
 ):
     if disable_instrumentations is None:
         for idx, (name, v) in enumerate(all_instrumentations.items()):
+
             v.instrument()
     else:
 

langtrace_python_sdk/utils/__init__.py

@@ -1,10 +1,21 @@
+from openai import NOT_GIVEN
 from .sdk_version_checker import SDKVersionChecker
+from opentelemetry.trace import Span
+from langtrace.trace_attributes import SpanAttributes
 
 
-def set_span_attribute(span, name, value):
+def set_span_attribute(span: Span, name, value):
     if value is not None:
-        if value != "":
-            span.set_attribute(name, value)
+        if value != "" or value != NOT_GIVEN:
+            if name == SpanAttributes.LLM_PROMPTS:
+                span.add_event(
+                    name=SpanAttributes.LLM_CONTENT_PROMPT,
+                    attributes={
+                        SpanAttributes.LLM_PROMPTS: value,
+                    },
+                )
+            else:
+                span.set_attribute(name, value)
     return
 
 

langtrace_python_sdk/utils/llm.py

@@ -14,10 +14,22 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
+from langtrace_python_sdk.utils import set_span_attribute
+from openai import NOT_GIVEN
 from tiktoken import get_encoding
 
-from langtrace_python_sdk.constants.instrumentation.common import TIKTOKEN_MODEL_MAPPING
+from langtrace_python_sdk.constants.instrumentation.common import (
+    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
+    TIKTOKEN_MODEL_MAPPING,
+)
 from langtrace_python_sdk.constants.instrumentation.openai import OPENAI_COST_TABLE
+from langtrace.trace_attributes import SpanAttributes, Event
+from importlib_metadata import version as v
+import json
+from opentelemetry import baggage
+from opentelemetry.trace import Span
+from opentelemetry.trace.status import StatusCode
 
 
 def estimate_tokens(prompt):
@@ -29,6 +41,15 @@ def estimate_tokens(prompt):
     return 0
 
 
+def set_event_completion_chunk(span: Span, chunk):
+    span.add_event(
+        name=SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK,
+        attributes={
+            SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: json.dumps(chunk),
+        },
+    )
+
+
 def estimate_tokens_using_tiktoken(prompt, model):
     """
     Estimate the number of tokens in a prompt using tiktoken."""
@@ -60,11 +81,110 @@ def calculate_price_from_usage(model, usage):
     return 0
 
 
-def set_span_attributes(span, name, value):
-    if value is not None:
-        if value != "":
-            span.set_attribute(name, value)
-    return
+def get_langtrace_attributes(version, service_provider, vendor_type="llm"):
+    return {
+        SpanAttributes.LANGTRACE_SDK_NAME: LANGTRACE_SDK_NAME,
+        SpanAttributes.LANGTRACE_VERSION: v(LANGTRACE_SDK_NAME),
+        SpanAttributes.LANGTRACE_SERVICE_VERSION: version,
+        SpanAttributes.LANGTRACE_SERVICE_NAME: service_provider,
+        SpanAttributes.LANGTRACE_SERVICE_TYPE: vendor_type,
+    }
+
+
+def get_llm_request_attributes(kwargs, prompts=None, model=None):
+
+    user = kwargs.get("user", None)
+    if prompts is None:
+        prompts = (
+            [{"role": user or "user", "content": kwargs.get("prompt")}]
+            if "prompt" in kwargs
+            else None
+        )
+    top_k = (
+        kwargs.get("n", None)
+        or kwargs.get("k", None)
+        or kwargs.get("top_k", None)
+        or kwargs.get("top_n", None)
+    )
+
+    top_p = kwargs.get("p", None) or kwargs.get("top_p", None)
+    tools = kwargs.get("tools", None)
+    return {
+        SpanAttributes.LLM_REQUEST_MODEL: model or kwargs.get("model"),
+        SpanAttributes.LLM_IS_STREAMING: kwargs.get("stream"),
+        SpanAttributes.LLM_REQUEST_TEMPERATURE: kwargs.get("temperature"),
+        SpanAttributes.LLM_TOP_K: top_k,
+        SpanAttributes.LLM_PROMPTS: json.dumps(prompts) if prompts else None,
+        SpanAttributes.LLM_USER: user,
+        SpanAttributes.LLM_REQUEST_TOP_P: top_p,
+        SpanAttributes.LLM_REQUEST_MAX_TOKENS: kwargs.get("max_tokens"),
+        SpanAttributes.LLM_SYSTEM_FINGERPRINT: kwargs.get("system_fingerprint"),
+        SpanAttributes.LLM_PRESENCE_PENALTY: kwargs.get("presence_penalty"),
+        SpanAttributes.LLM_FREQUENCY_PENALTY: kwargs.get("frequency_penalty"),
+        SpanAttributes.LLM_REQUEST_SEED: kwargs.get("seed"),
+        SpanAttributes.LLM_TOOLS: json.dumps(tools) if tools else None,
+        SpanAttributes.LLM_REQUEST_LOGPROPS: kwargs.get("logprobs"),
+        SpanAttributes.LLM_REQUEST_LOGITBIAS: kwargs.get("logit_bias"),
+        SpanAttributes.LLM_REQUEST_TOP_LOGPROPS: kwargs.get("top_logprobs"),
+    }
+
+
+def get_extra_attributes():
+    extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
+    return extra_attributes or {}
+
+
+def get_llm_url(instance):
+    return {
+        SpanAttributes.LLM_URL: get_base_url(instance),
+    }
+
+
+def get_base_url(instance):
+    return (
+        str(instance._client._base_url)
+        if hasattr(instance, "_client") and hasattr(instance._client, "_base_url")
+        else ""
+    )
+
+
+def is_streaming(kwargs):
+    return not (
+        kwargs.get("stream") is False
+        or kwargs.get("stream") is None
+        or kwargs.get("stream") == NOT_GIVEN
+    )
+
+
+def set_usage_attributes(span, usage):
+    if usage is None:
+        return
+
+    input_tokens = usage.get("input_tokens") or usage.get("prompt_tokens") or 0
+    output_tokens = usage.get("output_tokens") or usage.get("completion_tokens") or 0
+
+    set_span_attribute(
+        span,
+        SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
+        input_tokens,
+    )
+
+    set_span_attribute(
+        span,
+        SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
+        output_tokens,
+    )
+
+    set_span_attribute(
+        span,
+        SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
+        input_tokens + output_tokens,
+    )
+
+    if "search_units" in usage:
+        set_span_attribute(
+            span, SpanAttributes.LLM_USAGE_SEARCH_UNITS, usage["search_units"]
+        )
 
 
 def get_tool_calls(item):
@@ -77,3 +197,188 @@ def get_tool_calls(item):
         if hasattr(item, "tool_calls") and item.tool_calls is not None:
             return item.tool_calls
         return None
+
+
+def set_event_completion(span: Span, result_content):
+
+    span.add_event(
+        name=SpanAttributes.LLM_CONTENT_COMPLETION,
+        attributes={
+            SpanAttributes.LLM_COMPLETIONS: json.dumps(result_content),
+        },
+    )
+
+
+def set_span_attributes(span: Span, attributes: dict):
+    for field, value in attributes.model_dump(by_alias=True).items():
+        set_span_attribute(span, field, value)
+
+
+class StreamWrapper:
+    span: Span
+
+    def __init__(
+        self, stream, span, prompt_tokens, function_call=False, tool_calls=False
+    ):
+        self.stream = stream
+        self.span = span
+        self.prompt_tokens = prompt_tokens
+        self.function_call = function_call
+        self.tool_calls = tool_calls
+        self.result_content = []
+        self.completion_tokens = 0
+        self._span_started = False
+        self.setup()
+
+    def setup(self):
+        if not self._span_started:
+            self.span.add_event(Event.STREAM_START.value)
+            self._span_started = True
+
+    def cleanup(self):
+        if self._span_started:
+            self.span.add_event(Event.STREAM_END.value)
+            set_span_attribute(
+                self.span,
+                SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
+                self.prompt_tokens,
+            )
+            set_span_attribute(
+                self.span,
+                SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
+                self.completion_tokens,
+            )
+            set_span_attribute(
+                self.span,
+                SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
+                self.prompt_tokens + self.completion_tokens,
+            )
+            set_event_completion(
+                self.span,
+                [
+                    {
+                        "role": "assistant",
+                        "content": "".join(self.result_content),
+                    }
+                ],
+            )
+
+            self.span.set_status(StatusCode.OK)
+            self.span.end()
+            self._span_started = False
+
+    def __enter__(self):
+        self.setup()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.cleanup()
+
+    async def __aenter__(self):
+        self.setup()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self.cleanup()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            chunk = next(self.stream)
+            self.process_chunk(chunk)
+            return chunk
+        except StopIteration:
+            self.cleanup()
+            raise
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            chunk = await self.stream.__anext__()
+            self.process_chunk(chunk)
+            return chunk
+        except StopAsyncIteration:
+            self.cleanup()
+            raise StopAsyncIteration
+
+    def process_chunk(self, chunk):
+        if hasattr(chunk, "model") and chunk.model is not None:
+            set_span_attribute(
+                self.span,
+                SpanAttributes.LLM_RESPONSE_MODEL,
+                chunk.model,
+            )
+
+        if hasattr(chunk, "choices") and chunk.choices is not None:
+            content = []
+            if not self.function_call and not self.tool_calls:
+                for choice in chunk.choices:
+                    if choice.delta and choice.delta.content is not None:
+                        token_counts = estimate_tokens(choice.delta.content)
+                        self.completion_tokens += token_counts
+                        content = [choice.delta.content]
+            elif self.function_call:
+                for choice in chunk.choices:
+                    if (
+                        choice.delta
+                        and choice.delta.function_call is not None
+                        and choice.delta.function_call.arguments is not None
+                    ):
+                        token_counts = estimate_tokens(
+                            choice.delta.function_call.arguments
+                        )
+                        self.completion_tokens += token_counts
+                        content = [choice.delta.function_call.arguments]
+            elif self.tool_calls:
+                for choice in chunk.choices:
+                    if choice.delta and choice.delta.tool_calls is not None:
+                        toolcalls = choice.delta.tool_calls
+                        content = []
+                        for tool_call in toolcalls:
+                            if (
+                                tool_call
+                                and tool_call.function is not None
+                                and tool_call.function.arguments is not None
+                            ):
+                                token_counts = estimate_tokens(
+                                    tool_call.function.arguments
+                                )
+                                self.completion_tokens += token_counts
+                                content.append(tool_call.function.arguments)
+            self.span.add_event(
+                Event.STREAM_OUTPUT.value,
+                {
+                    SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: (
+                        "".join(content)
+                        if len(content) > 0 and content[0] is not None
+                        else ""
+                    )
+                },
+            )
+            if content:
+                self.result_content.append(content[0])
+
+        if hasattr(chunk, "text"):
+            token_counts = estimate_tokens(chunk.text)
+            self.completion_tokens += token_counts
+            content = [chunk.text]
+            self.span.add_event(
+                Event.STREAM_OUTPUT.value,
+                {
+                    SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: (
+                        "".join(content)
+                        if len(content) > 0 and content[0] is not None
+                        else ""
+                    )
+                },
+            )
+            if content:
+                self.result_content.append(content[0])
+
+        if hasattr(chunk, "usage_metadata"):
+            self.completion_tokens = chunk.usage_metadata.candidates_token_count
+            self.prompt_tokens = chunk.usage_metadata.prompt_token_count

langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "2.1.28"
+__version__ = "2.2.1"