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

langtrace_python_sdk/instrumentation/gemini/patch.py

@@ -0,0 +1,186 @@
+from langtrace.trace_attributes import LLMSpanAttributes, SpanAttributes
+from opentelemetry import trace
+from opentelemetry.trace import Span, SpanKind, Tracer
+from opentelemetry.trace.propagation import set_span_in_context
+from opentelemetry.trace.status import Status, StatusCode
+
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
+from langtrace_python_sdk.utils.llm import (
+    get_extra_attributes,
+    get_langtrace_attributes,
+    get_llm_request_attributes,
+    get_llm_url,
+    is_streaming,
+    set_event_completion,
+    set_event_completion_chunk,
+    set_span_attributes,
+    set_usage_attributes,
+)
+
+
+def patch_gemini(name, version, tracer: Tracer):
+    def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS["GEMINI"]
+        prompts = serialize_prompts(args, kwargs, instance)
+        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(kwargs):
+                return build_streaming_response(span, result)
+
+            else:
+                set_response_attributes(span, result)
+            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
+
+
+def apatch_gemini(name, version, tracer: Tracer):
+    async def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS["GEMINI"]
+        prompts = serialize_prompts(args, kwargs, instance)
+        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 = await wrapped(*args, **kwargs)
+            if is_streaming(kwargs):
+                return abuild_streaming_response(span, result)
+            else:
+                set_response_attributes(span, result)
+            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
+
+
+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("models/", "")
+    return llm_model
+
+
+def serialize_prompts(args, kwargs, instance):
+    prompts = []
+    if hasattr(instance, "_system_instruction") and instance._system_instruction is not None:
+        system_prompt = {
+            "role": "system",
+            "content": instance._system_instruction.__dict__["_pb"].parts[0].text,
+        }
+        prompts.append(system_prompt)
+
+    if args is not None and len(args) > 0:
+        content = ""
+        for arg in args:
+            if isinstance(arg, str):
+                content = f"{content}{arg}\n"
+            elif isinstance(arg, list):
+                for subarg in arg:
+                    content = f"{content}{subarg}\n"
+        prompts.append({"role": "user", "content": content})
+    return prompts
+
+
+def set_response_attributes(
+    span: Span,
+    result,
+):
+    span.set_status(Status(StatusCode.OK))
+    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}
+        )
+
+
+def build_streaming_response(span, response):
+    complete_response = ""
+    for item in response:
+        item_to_yield = item
+        complete_response += str(item.text)
+        yield item_to_yield
+        set_event_completion_chunk(span, item.text)
+        if hasattr(item, "usage_metadata"):
+            usage = item.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}
+            )
+
+    set_response_attributes(span, response)
+    span.set_status(Status(StatusCode.OK))
+    span.end()
+
+
+async def abuild_streaming_response(span, response):
+    complete_response = ""
+    async for item in response:
+        item_to_yield = item
+        complete_response += str(item.text)
+        yield item_to_yield
+        set_event_completion_chunk(span, item.text)
+        if hasattr(item, "usage_metadata"):
+            usage = item.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}
+            )
+
+    set_response_attributes(span, response)
+    span.set_status(Status(StatusCode.OK))
+    span.end()

langtrace_python_sdk/instrumentation/groq/patch.py

@@ -17,11 +17,21 @@ limitations under the License.
 import json
 
 from langtrace.trace_attributes import Event, LLMSpanAttributes
+from langtrace_python_sdk.utils import set_span_attribute
 from opentelemetry import baggage, trace
 from opentelemetry.trace.propagation import set_span_in_context
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
 
+from langtrace_python_sdk.utils.llm import (
+    get_base_url,
+    get_extra_attributes,
+    get_llm_request_attributes,
+    get_llm_url,
+    get_langtrace_attributes,
+    set_event_completion,
+    set_usage_attributes,
+)
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
     SERVICE_PROVIDERS,
@@ -31,26 +41,20 @@ from langtrace_python_sdk.utils.llm import calculate_prompt_tokens, estimate_tok
 from importlib_metadata import version as v
 
 from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
+from langtrace.trace_attributes import SpanAttributes
 
 
 def chat_completions_create(original_method, version, tracer):
     """Wrap the `create` method of the `ChatCompletion` class to trace it."""
 
     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["GROQ"]
         # If base url contains perplexity or azure, set the service provider accordingly
-        if "perplexity" in base_url:
+        if "perplexity" in get_base_url(instance):
             service_provider = SERVICE_PROVIDERS["PPLX"]
-        elif "azure" in base_url:
+        elif "azure" in get_base_url(instance):
             service_provider = SERVICE_PROVIDERS["AZURE"]
 
-        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
-
         # handle tool calls in the kwargs
         llm_prompts = []
         for item in kwargs.get("messages", []):
@@ -80,27 +84,16 @@ def chat_completions_create(original_method, version, tracer):
                 llm_prompts.append(item)
 
         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),
-            "url.full": base_url,
-            "llm.api": APIS["CHAT_COMPLETION"]["ENDPOINT"],
-            "llm.prompts": json.dumps(llm_prompts),
-            "llm.stream": kwargs.get("stream"),
-            **(extra_attributes if extra_attributes is not None else {}),
+            **get_langtrace_attributes(version, service_provider),
+            **get_llm_request_attributes(kwargs, prompts=llm_prompts),
+            **get_llm_url(instance),
+            SpanAttributes.LLM_PATH: APIS["CHAT_COMPLETION"]["ENDPOINT"],
+            **get_extra_attributes(),
         }
 
         attributes = LLMSpanAttributes(**span_attributes)
 
         tools = []
-        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")
         if kwargs.get("functions") is not None:
             for function in kwargs.get("functions"):
                 tools.append(json.dumps({"type": "function", "function": function}))
@@ -111,20 +104,21 @@ def chat_completions_create(original_method, version, tracer):
 
         # TODO(Karthik): Gotta figure out how to handle streaming with context
         # with tracer.start_as_current_span(APIS["CHAT_COMPLETION"]["METHOD"],
-        #                                   kind=SpanKind.CLIENT) as span:
+        #                                   kind=SpanKind.CLIENT.value) as span:
         span = tracer.start_span(
             APIS["CHAT_COMPLETION"]["METHOD"],
-            kind=SpanKind.CLIENT,
+            kind=SpanKind.CLIENT.value,
             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:
-                span.set_attribute(field, value)
+            set_span_attribute(span, field, value)
         try:
             # Attempt to call the original method
             result = wrapped(*args, **kwargs)
             if kwargs.get("stream") is False or kwargs.get("stream") is None:
-                span.set_attribute("llm.model", result.model)
+                set_span_attribute(
+                    span, SpanAttributes.LLM_RESPONSE_MODEL, result.model
+                )
                 if hasattr(result, "choices") and result.choices is not None:
                     responses = [
                         {
@@ -146,27 +140,23 @@ def chat_completions_create(original_method, version, tracer):
                         }
                         for choice in result.choices
                     ]
-                    span.set_attribute("llm.responses", json.dumps(responses))
-                else:
-                    responses = []
-                    span.set_attribute("llm.responses", json.dumps(responses))
+                    set_event_completion(span, responses)
+
                 if (
                     hasattr(result, "system_fingerprint")
                     and result.system_fingerprint is not None
                 ):
-                    span.set_attribute(
-                        "llm.system.fingerprint", result.system_fingerprint
+                    set_span_attribute(
+                        span,
+                        SpanAttributes.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 = {
-                            "input_tokens": result.usage.prompt_tokens,
-                            "output_tokens": usage.completion_tokens,
-                            "total_tokens": usage.total_tokens,
-                        }
-                        span.set_attribute("llm.token.counts", json.dumps(usage_dict))
+                    set_usage_attributes(span, dict(usage))
+
                 span.set_status(StatusCode.OK)
                 span.end()
                 return result
@@ -255,7 +245,7 @@ def chat_completions_create(original_method, version, tracer):
                 span.add_event(
                     Event.STREAM_OUTPUT.value,
                     {
-                        "response": (
+                        SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: (
                             "".join(content)
                             if len(content) > 0 and content[0] is not None
                             else ""
@@ -267,27 +257,14 @@ def chat_completions_create(original_method, version, tracer):
         finally:
             # Finalize span after processing all chunks
             span.add_event(Event.STREAM_END.value)
-            span.set_attribute(
-                "llm.token.counts",
-                json.dumps(
-                    {
-                        "input_tokens": prompt_tokens,
-                        "output_tokens": completion_tokens,
-                        "total_tokens": prompt_tokens + completion_tokens,
-                    }
-                ),
+            set_usage_attributes(
+                span,
+                {"input_tokens": prompt_tokens, "output_tokens": completion_tokens},
             )
-            span.set_attribute(
-                "llm.responses",
-                json.dumps(
-                    [
-                        {
-                            "role": "assistant",
-                            "content": "".join(result_content),
-                        }
-                    ]
-                ),
+            set_event_completion(
+                span, [{"role": "assistant", "content": "".join(result_content)}]
             )
+
             span.set_status(StatusCode.OK)
             span.end()
 
@@ -299,20 +276,13 @@ def async_chat_completions_create(original_method, version, tracer):
     """Wrap the `create` method of the `ChatCompletion` class to trace it."""
 
     async 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["GROQ"]
         # If base url contains perplexity or azure, set the service provider accordingly
-        if "perplexity" in base_url:
+        if "perplexity" in get_base_url(instance):
             service_provider = SERVICE_PROVIDERS["PPLX"]
-        elif "azure" in base_url:
+        elif "azure" in get_base_url(instance):
             service_provider = SERVICE_PROVIDERS["AZURE"]
 
-        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
-
         # handle tool calls in the kwargs
         llm_prompts = []
         for item in kwargs.get("messages", []):
@@ -342,27 +312,17 @@ def async_chat_completions_create(original_method, version, tracer):
                 llm_prompts.append(item)
 
         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),
-            "url.full": base_url,
-            "llm.api": APIS["CHAT_COMPLETION"]["ENDPOINT"],
-            "llm.prompts": json.dumps(llm_prompts),
-            "llm.stream": kwargs.get("stream"),
-            **(extra_attributes if extra_attributes is not None else {}),
+            **get_langtrace_attributes(version, service_provider),
+            **get_llm_request_attributes(kwargs, prompts=llm_prompts),
+            **get_llm_url(instance),
+            SpanAttributes.LLM_PATH: APIS["CHAT_COMPLETION"]["ENDPOINT"],
+            **get_extra_attributes(),
         }
 
         attributes = LLMSpanAttributes(**span_attributes)
 
         tools = []
-        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")
+
         if kwargs.get("functions") is not None:
             for function in kwargs.get("functions"):
                 tools.append(json.dumps({"type": "function", "function": function}))
@@ -373,18 +333,19 @@ def async_chat_completions_create(original_method, version, tracer):
 
         # TODO(Karthik): Gotta figure out how to handle streaming with context
         # with tracer.start_as_current_span(APIS["CHAT_COMPLETION"]["METHOD"],
-        #                                   kind=SpanKind.CLIENT) as span:
+        #                                   kind=SpanKind.CLIENT.value) as span:
         span = tracer.start_span(
-            APIS["CHAT_COMPLETION"]["METHOD"], kind=SpanKind.CLIENT
+            APIS["CHAT_COMPLETION"]["METHOD"], kind=SpanKind.CLIENT.value
         )
         for field, value in attributes.model_dump(by_alias=True).items():
-            if value is not None:
-                span.set_attribute(field, value)
+            set_span_attribute(span, field, value)
         try:
             # Attempt to call the original method
             result = await wrapped(*args, **kwargs)
             if kwargs.get("stream") is False or kwargs.get("stream") is None:
-                span.set_attribute("llm.model", result.model)
+                set_span_attribute(
+                    span, SpanAttributes.LLM_RESPONSE_MODEL, result.model
+                )
                 if hasattr(result, "choices") and result.choices is not None:
                     responses = [
                         {
@@ -406,27 +367,25 @@ def async_chat_completions_create(original_method, version, tracer):
                         }
                         for choice in result.choices
                     ]
-                    span.set_attribute("llm.responses", json.dumps(responses))
-                else:
-                    responses = []
-                    span.set_attribute("llm.responses", json.dumps(responses))
+
+                    set_event_completion(span, responses)
+
                 if (
                     hasattr(result, "system_fingerprint")
                     and result.system_fingerprint is not None
                 ):
-                    span.set_attribute(
-                        "llm.system.fingerprint", result.system_fingerprint
+                    set_span_attribute(
+                        span,
+                        SpanAttributes.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 = {
-                            "input_tokens": result.usage.prompt_tokens,
-                            "output_tokens": usage.completion_tokens,
-                            "total_tokens": usage.total_tokens,
-                        }
-                        span.set_attribute("llm.token.counts", json.dumps(usage_dict))
+                        set_usage_attributes(span, dict(usage))
+
                 span.set_status(StatusCode.OK)
                 span.end()
                 return result
@@ -469,6 +428,9 @@ def async_chat_completions_create(original_method, version, tracer):
         try:
             async for chunk in result:
                 if hasattr(chunk, "model") and chunk.model is not None:
+                    set_span_attribute(
+                        span, SpanAttributes.LLM_RESPONSE_MODEL, chunk.model
+                    )
                     span.set_attribute("llm.model", chunk.model)
                 if hasattr(chunk, "choices") and chunk.choices is not None:
                     if not function_call and not tool_calls:
@@ -513,9 +475,9 @@ def async_chat_completions_create(original_method, version, tracer):
                 else:
                     content = []
                 span.add_event(
-                    Event.STREAM_OUTPUT.value,
+                    Event.RESPONSE.value,
                     {
-                        "response": (
+                        SpanAttributes.LLM_COMPLETIONS: (
                             "".join(content)
                             if len(content) > 0 and content[0] is not None
                             else ""
@@ -527,27 +489,22 @@ def async_chat_completions_create(original_method, version, tracer):
         finally:
             # Finalize span after processing all chunks
             span.add_event(Event.STREAM_END.value)
-            span.set_attribute(
-                "llm.token.counts",
-                json.dumps(
+
+            set_usage_attributes(
+                span,
+                {"input_tokens": prompt_tokens, "output_tokens": completion_tokens},
+            )
+
+            set_event_completion(
+                span,
+                [
                     {
-                        "input_tokens": prompt_tokens,
-                        "output_tokens": completion_tokens,
-                        "total_tokens": prompt_tokens + completion_tokens,
+                        "role": "assistant",
+                        "content": "".join(result_content),
                     }
-                ),
-            )
-            span.set_attribute(
-                "llm.responses",
-                json.dumps(
-                    [
-                        {
-                            "role": "assistant",
-                            "content": "".join(result_content),
-                        }
-                    ]
-                ),
+                ],
             )
+
             span.set_status(StatusCode.OK)
             span.end()