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

tests/utils.py

@@ -1,5 +1,8 @@
 from unittest.mock import MagicMock, patch
 import json
+from langtrace.trace_attributes import SpanAttributes
+from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
+from importlib_metadata import version as v
 
 
 def common_setup(data, method_to_mock=None):
@@ -22,10 +25,9 @@ def common_setup(data, method_to_mock=None):
 
 
 def assert_token_count(attributes):
-    tokens = json.loads(attributes.get("llm.token.counts"))
-    output_tokens = tokens.get("output_tokens")
-    prompt_tokens = tokens.get("input_tokens")
-    total_tokens = tokens.get("total_tokens")
+    output_tokens = attributes.get(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS)
+    prompt_tokens = attributes.get(SpanAttributes.LLM_USAGE_PROMPT_TOKENS)
+    total_tokens = attributes.get(SpanAttributes.LLM_USAGE_TOTAL_TOKENS)
 
     assert (
         output_tokens is not None
@@ -36,9 +38,42 @@ def assert_token_count(attributes):
 
 
 def assert_response_format(attributes):
-    langtrace_responses = json.loads(attributes.get("llm.responses"))
+
+    langtrace_responses = json.loads(attributes.get(SpanAttributes.LLM_COMPLETIONS))
+
     assert isinstance(langtrace_responses, list)
     for langtrace_response in langtrace_responses:
         assert isinstance(langtrace_response, dict)
         assert "role" in langtrace_response
         assert "content" in langtrace_response
+
+
+def assert_langtrace_attributes(attributes, vendor, vendor_type="llm"):
+
+    assert attributes.get(SpanAttributes.LANGTRACE_SDK_NAME) == LANGTRACE_SDK_NAME
+    assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_NAME) == vendor
+    assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_TYPE) == vendor_type
+    assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_VERSION) == v(vendor.lower())
+    assert attributes.get(SpanAttributes.LANGTRACE_VERSION) == v(LANGTRACE_SDK_NAME)
+
+
+def assert_prompt_in_events(
+    events,
+):
+    prompt_event = list(
+        filter(lambda event: event.name == SpanAttributes.LLM_CONTENT_PROMPT, events)
+    )
+
+    assert prompt_event
+
+
+def assert_completion_in_events(
+    events,
+):
+    completion_event = list(
+        filter(
+            lambda event: event.name == SpanAttributes.LLM_CONTENT_COMPLETION, events
+        )
+    )
+
+    assert completion_event