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

langtrace_python_sdk/instrumentation/anthropic/patch.py

@@ -17,120 +17,59 @@ limitations under the License.
 import json
 
 from langtrace.trace_attributes import Event, LLMSpanAttributes
-from opentelemetry import baggage
+from langtrace_python_sdk.utils import set_span_attribute, silently_fail
+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_usage_attributes,
+)
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
+from langtrace.trace_attributes import SpanAttributes
 
 from langtrace_python_sdk.constants.instrumentation.anthropic import APIS
 from langtrace_python_sdk.constants.instrumentation.common import (
-    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
     SERVICE_PROVIDERS,
 )
-from importlib_metadata import version as v
-
-from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
 
 
 def messages_create(original_method, version, tracer):
     """Wrap the `messages_create` method."""
 
     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["ANTHROPIC"]
 
         # extract system from kwargs and attach as a role to the prompts
         # we do this to keep it consistent with the openai
-        prompts = json.dumps(kwargs.get("messages", []))
+        prompts = kwargs.get("messages", [])
         system = kwargs.get("system")
         if system:
-            prompts = json.dumps(
-                [{"role": "system", "content": system}] + kwargs.get("messages", [])
-            )
-        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
+            prompts = [{"role": "system", "content": system}] + kwargs.get("messages", [])
 
         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["MESSAGES_CREATE"]["ENDPOINT"],
-            "llm.model": kwargs.get("model"),
-            "llm.prompts": 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=prompts),
+            **get_llm_url(instance),
+            SpanAttributes.LLM_PATH: APIS["MESSAGES_CREATE"]["ENDPOINT"],
+            **get_extra_attributes(),
         }
 
         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("top_k") is not None:
-            attributes.llm_top_p = kwargs.get("top_k")
-        if kwargs.get("user") is not None:
-            attributes.llm_user = kwargs.get("user")
-        if kwargs.get("max_tokens") is not None:
-            attributes.llm_max_tokens = str(kwargs.get("max_tokens"))
-
         span = tracer.start_span(
             APIS["MESSAGES_CREATE"]["METHOD"], kind=SpanKind.CLIENT
         )
         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:
-                if hasattr(result, "content") and result.content is not None:
-                    span.set_attribute(
-                        "llm.model",
-                        result.model if result.model else kwargs.get("model"),
-                    )
-                    span.set_attribute(
-                        "llm.responses",
-                        json.dumps(
-                            [
-                                {
-                                    "role": result.role if result.role else "assistant",
-                                    "content": result.content[0].text,
-                                    "type": result.content[0].type,
-                                }
-                            ]
-                        ),
-                    )
-                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 = {
-                            "input_tokens": usage.input_tokens,
-                            "output_tokens": usage.output_tokens,
-                            "total_tokens": usage.input_tokens + usage.output_tokens,
-                        }
-                        span.set_attribute("llm.token.counts", json.dumps(usage_dict))
-                span.set_status(StatusCode.OK)
-                span.end()
-                return result
-            else:
-                return handle_streaming_response(result, span)
+            return set_response_attributes(result, span, kwargs)
+
         except Exception as err:
             # Record the exception in the span
             span.record_exception(err)
@@ -154,7 +93,9 @@ def messages_create(original_method, version, tracer):
                     and hasattr(chunk.message, "model")
                     and chunk.message.model is not None
                 ):
-                    span.set_attribute("llm.model", chunk.message.model)
+                    span.set_attribute(
+                        SpanAttributes.LLM_RESPONSE_MODEL, chunk.message.model
+                    )
                 content = ""
                 if hasattr(chunk, "delta") and chunk.delta is not None:
                     content = chunk.delta.text if hasattr(chunk.delta, "text") else ""
@@ -177,7 +118,8 @@ def messages_create(original_method, version, tracer):
                 # Add event for each chunk of content
                 if content:
                     span.add_event(
-                        Event.STREAM_OUTPUT.value, {"response": "".join(content)}
+                        Event.STREAM_OUTPUT.value,
+                        {SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: "".join(content)},
                     )
 
                 # Assuming this is part of a generator, yield chunk or aggregated content
@@ -186,22 +128,52 @@ def messages_create(original_method, version, tracer):
 
             # 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": input_tokens, "output_tokens": output_tokens}
+            )
+            completion = [{"role": "assistant", "content": "".join(result_content)}]
+            set_event_completion(span, completion)
+
+            span.set_status(StatusCode.OK)
+            span.end()
+
+    def set_response_attributes(result, span, kwargs):
+        if not is_streaming(kwargs):
+            if hasattr(result, "content") and result.content is not None:
+                set_span_attribute(
+                    span, SpanAttributes.LLM_RESPONSE_MODEL, result.model
+                )
+                completion = [
                     {
-                        "input_tokens": input_tokens,
-                        "output_tokens": output_tokens,
-                        "total_tokens": input_tokens + output_tokens,
+                        "role": result.role if result.role else "assistant",
+                        "content": result.content[0].text,
+                        "type": result.content[0].type,
                     }
-                ),
-            )
-            span.set_attribute(
-                "llm.responses",
-                json.dumps([{"role": "assistant", "content": "".join(result_content)}]),
-            )
+                ]
+                set_event_completion(span, completion)
+
+            else:
+                responses = []
+                set_event_completion(span, responses)
+
+            if (
+                hasattr(result, "system_fingerprint")
+                and result.system_fingerprint is not None
+            ):
+                span.set_attribute(
+                    SpanAttributes.LLM_SYSTEM_FINGERPRINT,
+                    result.system_fingerprint,
+                )
+            # Get the usage
+            if hasattr(result, "usage") and result.usage is not None:
+                usage = result.usage
+                set_usage_attributes(span, dict(usage))
+
             span.set_status(StatusCode.OK)
             span.end()
+            return result
+        else:
+            return handle_streaming_response(result, span)
 
     # return the wrapped method
     return traced_method

langtrace_python_sdk/instrumentation/chroma/patch.py

@@ -15,7 +15,7 @@ limitations under the License.
 """
 
 from langtrace.trace_attributes import DatabaseSpanAttributes
-from langtrace_python_sdk.utils.llm import set_span_attributes
+from langtrace_python_sdk.utils import set_span_attribute
 from langtrace_python_sdk.utils.silently_fail import silently_fail
 from opentelemetry import baggage, trace
 from opentelemetry.trace import SpanKind
@@ -119,20 +119,20 @@ def handle_null_params(param):
 
 @silently_fail
 def _set_chroma_add_attributes(span, kwargs):
-    set_span_attributes(
+    set_span_attribute(
         span, "db.chroma.add.ids_count", get_count_or_none(kwargs.get("ids"))
     )
-    set_span_attributes(
+    set_span_attribute(
         span,
         "db.chroma.add.embeddings_count",
         get_count_or_none(kwargs.get("embeddings")),
     )
-    set_span_attributes(
+    set_span_attribute(
         span,
         "db.chroma.add.metadatas_count",
         get_count_or_none(kwargs.get("metadatas")),
     )
-    set_span_attributes(
+    set_span_attribute(
         span,
         "db.chroma.add.documents_count",
         get_count_or_none(kwargs.get("documents")),
@@ -141,71 +141,71 @@ def _set_chroma_add_attributes(span, kwargs):
 
 @silently_fail
 def _set_chroma_get_attributes(span, kwargs):
-    set_span_attributes(
+    set_span_attribute(
         span, "db.chroma.get.ids_count", get_count_or_none(kwargs.get("ids"))
     )
-    set_span_attributes(
+    set_span_attribute(
         span, "db.chroma.get.where", handle_null_params(kwargs.get("where"))
     )
-    set_span_attributes(span, "db.chroma.get.limit", kwargs.get("limit"))
-    set_span_attributes(span, "db.chroma.get.offset", kwargs.get("offset"))
-    set_span_attributes(
+    set_span_attribute(span, "db.chroma.get.limit", kwargs.get("limit"))
+    set_span_attribute(span, "db.chroma.get.offset", kwargs.get("offset"))
+    set_span_attribute(
         span,
         "db.chroma.get.where_document",
         handle_null_params(kwargs.get("where_document")),
     )
-    set_span_attributes(
+    set_span_attribute(
         span, "db.chroma.get.include", handle_null_params(kwargs.get("include"))
     )
 
 
 @silently_fail
 def _set_chroma_query_attributes(span, kwargs):
-    set_span_attributes(
+    set_span_attribute(
         span,
         "db.chroma.query.query_embeddings_count",
         get_count_or_none(kwargs.get("query_embeddings")),
     )
-    set_span_attributes(
+    set_span_attribute(
         span,
         "db.chroma.query.query_texts_count",
         get_count_or_none(kwargs.get("query_texts")),
     )
-    set_span_attributes(span, "db.chroma.query.n_results", kwargs.get("n_results"))
-    set_span_attributes(
+    set_span_attribute(span, "db.chroma.query.n_results", kwargs.get("n_results"))
+    set_span_attribute(
         span, "db.chroma.query.where", handle_null_params(kwargs.get("where"))
     )
-    set_span_attributes(
+    set_span_attribute(
         span,
         "db.chroma.query.where_document",
         handle_null_params(kwargs.get("where_document")),
     )
-    set_span_attributes(
+    set_span_attribute(
         span, "db.chroma.query.include", handle_null_params(kwargs.get("include"))
     )
 
 
 @silently_fail
 def _set_chroma_peek_attributes(span, kwargs):
-    set_span_attributes(span, "db.chroma.peek.limit", kwargs.get("limit"))
+    set_span_attribute(span, "db.chroma.peek.limit", kwargs.get("limit"))
 
 
 @silently_fail
 def _set_chroma_update_attributes(span, kwargs):
-    set_span_attributes(
+    set_span_attribute(
         span, "db.chroma.update.ids_count", get_count_or_none(kwargs.get("ids"))
     )
-    set_span_attributes(
+    set_span_attribute(
         span,
         "db.chroma.update.embeddings_count",
         get_count_or_none(kwargs.get("embeddings")),
     )
-    set_span_attributes(
+    set_span_attribute(
         span,
         "db.chroma.update.metadatas_count",
         get_count_or_none(kwargs.get("metadatas")),
     )
-    set_span_attributes(
+    set_span_attribute(
         span,
         "db.chroma.update.documents_count",
         get_count_or_none(kwargs.get("documents")),
@@ -214,23 +214,23 @@ def _set_chroma_update_attributes(span, kwargs):
 
 @silently_fail
 def _set_chroma_modify_attributes(span, kwargs):
-    set_span_attributes(span, "db.chroma.modify.name", kwargs.get("name"))
+    set_span_attribute(span, "db.chroma.modify.name", kwargs.get("name"))
     # TODO: Add metadata attribute
 
 
 @silently_fail
 def _set_chroma_upsert_attributes(span, kwargs):
-    set_span_attributes(
+    set_span_attribute(
         span,
         "db.chroma.upsert.embeddings_count",
         get_count_or_none(kwargs.get("embeddings")),
     )
-    set_span_attributes(
+    set_span_attribute(
         span,
         "db.chroma.upsert.metadatas_count",
         get_count_or_none(kwargs.get("metadatas")),
     )
-    set_span_attributes(
+    set_span_attribute(
         span,
         "db.chroma.upsert.documents_count",
         get_count_or_none(kwargs.get("documents")),
@@ -239,13 +239,13 @@ def _set_chroma_upsert_attributes(span, kwargs):
 
 @silently_fail
 def _set_chroma_delete_attributes(span, kwargs):
-    set_span_attributes(
+    set_span_attribute(
         span, "db.chroma.delete.ids_count", get_count_or_none(kwargs.get("ids"))
     )
-    set_span_attributes(
+    set_span_attribute(
         span, "db.chroma.delete.where", handle_null_params(kwargs.get("where"))
     )
-    set_span_attributes(
+    set_span_attribute(
         span,
         "db.chroma.delete.where_document",
         handle_null_params(kwargs.get("where_document")),