Merge branch 'development' of github.com:Scale3-Labs/langtrace-python-sdk into development

Author: karthikscale3

src/langtrace_python_sdk/instrumentation/milvus/patch.py

@@ -1,8 +1,10 @@
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
 from langtrace_python_sdk.utils.silently_fail import silently_fail
 from opentelemetry.trace import Tracer
 from opentelemetry.trace import SpanKind
 from langtrace_python_sdk.utils import handle_span_error, set_span_attribute
 from langtrace_python_sdk.utils.llm import (
+    get_langtrace_attributes,
     get_extra_attributes,
     set_span_attributes,
 )
@@ -16,6 +18,11 @@ def traced_method(wrapped, instance, args, kwargs):
         with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
             try:
                 span_attributes = {
+                    **get_langtrace_attributes(
+                        service_provider=SERVICE_PROVIDERS["MILVUS"],
+                        version=version,
+                        vendor_type="Vector Database",
+                    ),
                     "db.system": "milvus",
                     "db.operation": operation,
                     "db.name": kwargs.get("collection_name", None),

src/langtrace_python_sdk/instrumentation/mistral/patch.py

@@ -46,7 +46,7 @@
 
 
 def chat_complete(original_method, version, tracer, is_async=False, is_streaming=False):
-        
+
     def traced_method(wrapped, instance, args, kwargs):
         service_provider = SERVICE_PROVIDERS["MISTRAL"]
         llm_prompts = []
@@ -65,7 +65,6 @@ def traced_method(wrapped, instance, args, kwargs):
 
         attributes = LLMSpanAttributes(**span_attributes)
 
-
         span = tracer.start_span(
             name=get_span_name(APIS[api]["METHOD"]),
             kind=SpanKind.CLIENT,
@@ -87,7 +86,6 @@ def traced_method(wrapped, instance, args, kwargs):
                 span.set_status(StatusCode.OK)
                 span.end()
                 return result
-                
 
         except Exception as error:
             span.record_exception(error)
@@ -186,7 +184,6 @@ def _set_response_attributes(span, kwargs, result):
             for choice in result.choices
         ]
         set_event_completion(span, responses)
-
     # Get the usage
     if hasattr(result, "usage") and result.usage is not None:
-        set_usage_attributes(span, result.usage)
+        set_usage_attributes(span, dict(result.usage))

src/langtrace_python_sdk/instrumentation/ollama/patch.py

@@ -84,7 +84,7 @@ async def traced_method(wrapped, instance, args, kwargs):
         try:
             result = await wrapped(*args, **kwargs)
             if kwargs.get("stream"):
-                return StreamWrapper(span, result)
+                return StreamWrapper(result, span)
             else:
                 _set_response_attributes(span, result)
             span.end()

src/langtrace_python_sdk/utils/with_root_span.py

@@ -14,6 +14,7 @@
 limitations under the License.
 """
 
+import base64
 import asyncio
 import os
 from deprecated import deprecated
@@ -145,7 +146,9 @@ class SendUserFeedback:
     _langtrace_api_key: str
 
     def __init__(self):
-        self._langtrace_host = os.environ.get("LANGTRACE_API_HOST", LANGTRACE_REMOTE_URL)
+        self._langtrace_host = os.environ.get(
+            "LANGTRACE_API_HOST", LANGTRACE_REMOTE_URL
+        )
         # When the host is set to /api/trace, remove the /api/trace
         if self._langtrace_host.endswith("/api/trace"):
             self._langtrace_host = self._langtrace_host.replace("/api/trace", "")
@@ -162,14 +165,13 @@ def evaluate(self, data: EvaluationAPIData) -> None:
                 print(Fore.RESET)
                 return
 
-            # convert spanId and traceId to hexadecimals
-            span_hex_number = hex(int(data["spanId"], 10))[2:]  # Convert to hex and remove the '0x' prefix
-            formatted_span_hex_number = span_hex_number.zfill(16)  # Pad with zeros to 16 characters
-            data["spanId"] = f"0x{formatted_span_hex_number}"
+            # convert spanId and traceId to Base64
+            span_bytes = int(data["spanId"], 10).to_bytes(8, "big")
+            data["spanId"] = base64.b64encode(span_bytes).decode("ascii")
 
-            trace_hex_number = hex(int(data["traceId"], 10))[2:]  # Convert to hex and remove the '0x' prefix
-            formatted_trace_hex_number = trace_hex_number.zfill(32)  # Pad with zeros to 32 characters
-            data["traceId"] = f"0x{formatted_trace_hex_number}"
+            # Convert traceId to base64
+            trace_bytes = int(data["traceId"], 10).to_bytes(16, "big")
+            data["traceId"] = base64.b64encode(trace_bytes).decode("ascii")
 
             evaluation = self.get_evaluation(data["spanId"])
             headers = {"x-api-key": self._langtrace_api_key}

src/langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "3.3.6"
+__version__ = "3.3.9"