veadk-python 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

veadk/tracing/telemetry/attributes/extractors/llm_attributes_extractors.py

@@ -18,6 +18,7 @@ from veadk.tracing.telemetry.attributes.extractors.types import (
     ExtractorResponse,
     LLMAttributesParams,
 )
+from veadk.utils.misc import safe_json_serialize
 
 
 def llm_gen_ai_request_model(params: LLMAttributesParams) -> ExtractorResponse:
@@ -102,10 +103,11 @@ def llm_gen_ai_usage_cache_read_input_tokens(
 def llm_gen_ai_prompt(params: LLMAttributesParams) -> ExtractorResponse:
     # a part is a message
     messages: list[dict] = []
+    idx = 0
 
     for content in params.llm_request.contents:
         if content.parts:
-            for idx, part in enumerate(content.parts):
+            for part in content.parts:
                 message = {}
                 # text part
                 if part.text:
@@ -133,13 +135,25 @@ def llm_gen_ai_prompt(params: LLMAttributesParams) -> ExtractorResponse:
                         else "<unknown_function_name>"
                     )
                     message[f"gen_ai.prompt.{idx}.tool_calls.0.function.arguments"] = (
-                        json.dumps(part.function_call.args)
+                        safe_json_serialize(part.function_call.args)
                         if part.function_call.args
                         else json.dumps({})
                     )
+                # image
+                if part.inline_data:
+                    message[f"gen_ai.prompt.{idx}.type"] = "image_url"
+                    message[f"gen_ai.prompt.{idx}.image_url.name"] = (
+                        part.inline_data.display_name.split("/")[-1]
+                        if part.inline_data.display_name
+                        else "<unknown_image_name>"
+                    )
+                    message[f"gen_ai.prompt.{idx}.image_url.url"] = (
+                        part.inline_data.display_name
+                    )
 
                 if message:
                     messages.append(message)
+                    idx += 1
 
     return ExtractorResponse(content=messages)
 
@@ -168,7 +182,7 @@ def llm_gen_ai_completion(params: LLMAttributesParams) -> ExtractorResponse:
                     else "<unknown_function_name>"
                 )
                 message[f"gen_ai.completion.{idx}.tool_calls.0.function.arguments"] = (
-                    json.dumps(part.function_call.args)
+                    safe_json_serialize(part.function_call.args)
                     if part.function_call.args
                     else json.dumps({})
                 )
@@ -178,127 +192,241 @@ def llm_gen_ai_completion(params: LLMAttributesParams) -> ExtractorResponse:
     return ExtractorResponse(content=messages)
 
 
-def llm_gen_ai_is_streaming(params: LLMAttributesParams) -> ExtractorResponse:
-    # return params.llm_request.stream
-    return ExtractorResponse(content=None)
-
-
-def llm_gen_ai_operation_name(params: LLMAttributesParams) -> ExtractorResponse:
-    return ExtractorResponse(content="chat")
-
-
-def llm_gen_ai_system_message(params: LLMAttributesParams) -> ExtractorResponse:
-    event_attributes = {
-        "content": str(params.llm_request.config.system_instruction),
-        "role": "system",
-    }
-    return ExtractorResponse(type="event", content=event_attributes)
-
-
-def llm_gen_ai_user_message(params: LLMAttributesParams) -> ExtractorResponse:
-    # a content is a message
-    messages = []
-
-    for content in params.llm_request.contents:
-        if content.role == "user":
-            message_parts = []
-
-            if content.parts:
-                if len(content.parts) == 1:
-                    if content.parts[0].text:
-                        message_parts.append(
-                            {
-                                "role": content.role,
-                                "content": content.parts[0].text,
-                            }
-                        )
-                    elif content.parts[0].function_response:
-                        message_parts.append(
-                            {
-                                "role": content.role,
-                                "content": str(
-                                    content.parts[0].function_response.response
-                                ),
-                            }
-                        )
-                else:
-                    message_part = {"role": content.role}
+def llm_gen_ai_messages(params: LLMAttributesParams) -> ExtractorResponse:
+    events = []
+
+    # system message
+    events.append(
+        {
+            "gen_ai.system.message": {
+                "role": "system",
+                "content": str(params.llm_request.config.system_instruction),
+            }
+        }
+    )
+
+    # user, tool, and assistant message
+    if params.llm_request and params.llm_request.contents:
+        for content in params.llm_request.contents:
+            if content and content.parts:
+                # content.role == "user"
+                #   part.function_response -> gen_ai.tool.message
+                #   not part.function_response -> gen_ai.user.message
+                # content.role == "model" -> gen_ai.assistant.message
+                if content.role == "user":
+                    user_event = {}
+                    user_event["gen_ai.user.message"] = {"role": content.role}
                     for idx, part in enumerate(content.parts):
-                        # text part
-                        if part.text:
-                            message_part[f"parts.{idx}.type"] = "text"
-                            message_part[f"parts.{idx}.content"] = part.text
-                        # function response
                         if part.function_response:
-                            message_part[f"parts.{idx}.type"] = "function"
-                            message_part[f"parts.{idx}.content"] = str(
-                                part.function_response
+                            events.append(
+                                {
+                                    "gen_ai.tool.message": {
+                                        "role": "tool",
+                                        "id": part.function_response.id,
+                                        "content": safe_json_serialize(
+                                            part.function_response.response
+                                        ),
+                                    }
+                                }
                             )
+                        else:
+                            if part.text:
+                                if len(content.parts) == 1:
+                                    user_event["gen_ai.user.message"].update(
+                                        {"content": part.text}
+                                    )
+                                else:
+                                    user_event["gen_ai.user.message"].update(
+                                        {
+                                            f"parts.{idx}.type": "text",
+                                            f"parts.{idx}.text": part.text,
+                                        },
+                                    )
+                            if part.inline_data:
+                                if len(content.parts) == 1:
+                                    # TODO(qingliang)
+                                    pass
+                                else:
+                                    user_event["gen_ai.user.message"].update(
+                                        {
+                                            f"parts.{idx}.type": "image_url",
+                                            f"parts.{idx}.image_url.name": (
+                                                part.inline_data.display_name.split(
+                                                    "/"
+                                                )[-1]
+                                                if part.inline_data.display_name
+                                                else "<unknown_image_name>"
+                                            ),
+                                            f"parts.{idx}.image_url.url": (
+                                                part.inline_data.display_name
+                                                if part.inline_data.display_name
+                                                else "<unknown_image_url>"
+                                            ),
+                                        }
+                                    )
+                    # in case of only function response
+                    if len(user_event["gen_ai.user.message"].items()) > 1:
+                        events.append(user_event)
+                elif content.role == "model":
+                    event = {}
+                    event["gen_ai.assistant.message"] = {"role": content.role}
+                    for idx, part in enumerate(content.parts):
+                        if part.text:
+                            event["gen_ai.assistant.message"].update(
+                                {
+                                    f"parts.{idx}.type": "text",
+                                    f"parts.{idx}.text": part.text,
+                                }
+                            )
+                        if part.function_call:
+                            event["gen_ai.assistant.message"].update(
+                                {
+                                    "tool_calls.0.id": str(part.function_call.id),
+                                    "tool_calls.0.type": "function",
+                                    "tool_calls.0.function.name": part.function_call.name
+                                    if part.function_call.name
+                                    else "<unknown_function_name>",
+                                    "tool_calls.0.function.arguments": safe_json_serialize(
+                                        part.function_call.args
+                                    )
+                                    if part.function_call.args
+                                    else json.dumps({}),
+                                }
+                            )
+                    events.append(event)
 
-                    message_parts.append(message_part)
-
-            if message_parts:
-                messages.extend(message_parts)
-
-    return ExtractorResponse(type="event", content=messages)
+    return ExtractorResponse(type="event_list", content=events)
 
 
-def llm_gen_ai_assistant_message(params: LLMAttributesParams) -> ExtractorResponse:
-    # a content is a message
-    messages = []
+def llm_gen_ai_is_streaming(params: LLMAttributesParams) -> ExtractorResponse:
+    # return params.llm_request.stream
+    return ExtractorResponse(content=None)
 
-    # each part in each content we make it a message
-    # e.g. 2 contents and 3 parts each means 6 messages
-    for content in params.llm_request.contents:
-        if content.role == "model":
-            message_parts = []
-
-            # each part we make it a message
-            if content.parts:
-                # only one part
-                if len(content.parts) == 1:
-                    if content.parts[0].text:
-                        message_parts.append(
-                            {
-                                "role": content.role,
-                                "content": content.parts[0].text,
-                            }
-                        )
-                    elif content.parts[0].function_call:
-                        pass
-                # multiple parts
-                else:
-                    message_part = {"role": content.role}
 
-                    for idx, part in enumerate(content.parts):
-                        # parse content
-                        if part.text:
-                            message_part[f"parts.{idx}.type"] = "text"
-                            message_part[f"parts.{idx}.content"] = part.text
-                        # parse tool_calls
-                        if part.function_call:
-                            message_part["tool_calls.0.id"] = (
-                                part.function_call.id
-                                if part.function_call.id
-                                else "<unkown_function_call_id>"
-                            )
-                            message_part["tool_calls.0.type"] = "function"
-                            message_part["tool_calls.0.function.name"] = (
-                                part.function_call.name
-                                if part.function_call.name
-                                else "<unknown_function_name>"
-                            )
-                            message_part["tool_calls.0.function.arguments"] = (
-                                json.dumps(part.function_call.args)
-                                if part.function_call.args
-                                else json.dumps({})
-                            )
-                    message_parts.append(message_part)
+def llm_gen_ai_operation_name(params: LLMAttributesParams) -> ExtractorResponse:
+    return ExtractorResponse(content="chat")
 
-            if message_parts:
-                messages.extend(message_parts)
 
-    return ExtractorResponse(type="event", content=messages)
+# def llm_gen_ai_system_message(params: LLMAttributesParams) -> ExtractorResponse:
+#     event_attributes = {
+#         "content": str(params.llm_request.config.system_instruction),
+#         "role": "system",
+#     }
+#     return ExtractorResponse(type="event", content=event_attributes)
+
+
+# def llm_gen_ai_user_message(params: LLMAttributesParams) -> ExtractorResponse:
+#     # a content is a message
+#     messages = []
+
+#     for content in params.llm_request.contents:
+#         if content.role == "user":
+#             message_parts = []
+
+#             if content.parts:
+#                 if len(content.parts) == 1:
+#                     if content.parts[0].text:
+#                         message_parts.append(
+#                             {
+#                                 "role": content.role,
+#                                 "content": content.parts[0].text,
+#                             }
+#                         )
+#                     elif content.parts[0].function_response:
+#                         message_parts.append(
+#                             {
+#                                 "role": content.role,
+#                                 "content": str(
+#                                     content.parts[0].function_response.response
+#                                 ),
+#                             }
+#                         )
+#                 else:
+#                     message_part = {"role": content.role}
+#                     for idx, part in enumerate(content.parts):
+#                         # text part
+#                         if part.text:
+#                             message_part[f"parts.{idx}.type"] = "text"
+#                             message_part[f"parts.{idx}.content"] = part.text
+#                         # function response
+#                         if part.function_response:
+#                             message_part[f"parts.{idx}.type"] = "function"
+#                             message_part[f"parts.{idx}.content"] = str(
+#                                 part.function_response
+#                             )
+#                         if part.inline_data:
+#                             message_part[f"parts.{idx}.type"] = "image_url"
+#                             message_part[f"parts.{idx}.image_url.name"] = (
+#                                 part.inline_data.display_name.split("/")[-1]
+#                             )
+#                             message_part[f"parts.{idx}.image_url.url"] = (
+#                                 part.inline_data.display_name
+#                             )
+
+#                     message_parts.append(message_part)
+
+#             if message_parts:
+#                 messages.extend(message_parts)
+
+#     return ExtractorResponse(type="event", content=messages)
+
+
+# def llm_gen_ai_assistant_message(params: LLMAttributesParams) -> ExtractorResponse:
+#     # a content is a message
+#     messages = []
+
+#     # each part in each content we make it a message
+#     # e.g. 2 contents and 3 parts each means 6 messages
+#     for content in params.llm_request.contents:
+#         if content.role == "model":
+#             message_parts = []
+
+#             # each part we make it a message
+#             if content.parts:
+#                 # only one part
+#                 if len(content.parts) == 1:
+#                     if content.parts[0].text:
+#                         message_parts.append(
+#                             {
+#                                 "role": content.role,
+#                                 "content": content.parts[0].text,
+#                             }
+#                         )
+#                     elif content.parts[0].function_call:
+#                         pass
+#                 # multiple parts
+#                 else:
+#                     message_part = {"role": content.role}
+
+#                     for idx, part in enumerate(content.parts):
+#                         # parse content
+#                         if part.text:
+#                             message_part[f"parts.{idx}.type"] = "text"
+#                             message_part[f"parts.{idx}.content"] = part.text
+#                         # parse tool_calls
+#                         if part.function_call:
+#                             message_part["tool_calls.0.id"] = (
+#                                 part.function_call.id
+#                                 if part.function_call.id
+#                                 else "<unkown_function_call_id>"
+#                             )
+#                             message_part["tool_calls.0.type"] = "function"
+#                             message_part["tool_calls.0.function.name"] = (
+#                                 part.function_call.name
+#                                 if part.function_call.name
+#                                 else "<unknown_function_name>"
+#                             )
+#                             message_part["tool_calls.0.function.arguments"] = (
+#                                 safe_json_serialize(part.function_call.args)
+#                                 if part.function_call.args
+#                                 else json.dumps({})
+#                             )
+#                     message_parts.append(message_part)
+
+#             if message_parts:
+#                 messages.extend(message_parts)
+
+#     return ExtractorResponse(type="event", content=messages)
 
 
 def llm_gen_ai_choice(params: LLMAttributesParams) -> ExtractorResponse:
@@ -326,7 +454,7 @@ def llm_gen_ai_choice(params: LLMAttributesParams) -> ExtractorResponse:
                     else "<unknown_function_name>"
                 )
                 message["message.tool_calls.0.function.arguments"] = (
-                    json.dumps(part.function_call.args)
+                    safe_json_serialize(part.function_call.args)
                     if part.function_call.args
                     else json.dumps({})
                 )
@@ -351,7 +479,7 @@ def llm_gen_ai_choice(params: LLMAttributesParams) -> ExtractorResponse:
                         else "<unknown_function_name>"
                     )
                     message["message.tool_calls.0.function.arguments"] = (
-                        json.dumps(part.function_call.args)
+                        safe_json_serialize(part.function_call.args)
                         if part.function_call.args
                         else json.dumps({})
                     )
@@ -359,34 +487,89 @@ def llm_gen_ai_choice(params: LLMAttributesParams) -> ExtractorResponse:
     return ExtractorResponse(type="event", content=message)
 
 
+def llm_input_value(params: LLMAttributesParams) -> ExtractorResponse:
+    return ExtractorResponse(
+        content=str(params.llm_request.model_dump(exclude_none=True))
+    )
+
+
+def llm_output_value(params: LLMAttributesParams) -> ExtractorResponse:
+    return ExtractorResponse(
+        content=str(params.llm_response.model_dump(exclude_none=True))
+    )
+
+
+def llm_gen_ai_request_functions(params: LLMAttributesParams) -> ExtractorResponse:
+    functions = []
+
+    for idx, (tool_name, tool_instance) in enumerate(
+        params.llm_request.tools_dict.items()
+    ):
+        functions.append(
+            {
+                f"gen_ai.request.functions.{idx}.name": tool_instance.name,
+                f"gen_ai.request.functions.{idx}.description": tool_instance.description,
+                f"gen_ai.request.functions.{idx}.parameters": str(
+                    tool_instance._get_declaration().parameters.model_dump(  # type: ignore
+                        exclude_none=True
+                    )
+                    if tool_instance._get_declaration()
+                    and tool_instance._get_declaration().parameters  # type: ignore
+                    else {}
+                ),
+            }
+        )
+
+    return ExtractorResponse(content=functions)
+
+
 LLM_ATTRIBUTES = {
-    # ===== request attributes =====
+    # -> 1. attributes
+    # -> 1.1. request
     "gen_ai.request.model": llm_gen_ai_request_model,
     "gen_ai.request.type": llm_gen_ai_request_type,
     "gen_ai.request.max_tokens": llm_gen_ai_request_max_tokens,
     "gen_ai.request.temperature": llm_gen_ai_request_temperature,
     "gen_ai.request.top_p": llm_gen_ai_request_top_p,
-    # ===== response attributes =====
+    # CozeLoop required
+    "gen_ai.request.functions": llm_gen_ai_request_functions,
+    # -> 1.2. response
     "gen_ai.response.model": llm_gen_ai_response_model,
     "gen_ai.response.stop_reason": llm_gen_ai_response_stop_reason,
     "gen_ai.response.finish_reason": llm_gen_ai_response_finish_reason,
-    # ===== streaming =====
+    # -> 1.3. streaming
     "gen_ai.is_streaming": llm_gen_ai_is_streaming,
-    # ===== span type =====
+    # -> 1.4. span kind
     "gen_ai.operation.name": llm_gen_ai_operation_name,
-    # ===== inputs and outputs =====
-    # events
-    "gen_ai.system.message": llm_gen_ai_system_message,
-    "gen_ai.user.message": llm_gen_ai_user_message,
-    "gen_ai.assistant.message": llm_gen_ai_assistant_message,
-    "gen_ai.choice": llm_gen_ai_choice,
-    # attributes
+    # -> 1.5. inputs
     "gen_ai.prompt": llm_gen_ai_prompt,
+    # -> 1.6. outputs
     "gen_ai.completion": llm_gen_ai_completion,
-    # ===== usage =====
+    # -> 1.7. usage
     "gen_ai.usage.input_tokens": llm_gen_ai_usage_input_tokens,
     "gen_ai.usage.output_tokens": llm_gen_ai_usage_output_tokens,
     "gen_ai.usage.total_tokens": llm_gen_ai_usage_total_tokens,
     "gen_ai.usage.cache_creation_input_tokens": llm_gen_ai_usage_cache_creation_input_tokens,
     "gen_ai.usage.cache_read_input_tokens": llm_gen_ai_usage_cache_read_input_tokens,
+    # -> 2. events
+    # -> 2.1. inputs
+    # In order to adapt OpenTelemetry and CozeLoop rendering,
+    # and avoid error sequence of tool-call and too-response,
+    # we use `llm_gen_ai_messages` to upload system message, user message,
+    # and assistant message together.
+    # Correct sequence: system message, user message, tool message,
+    # and assistant message.
+    "gen_ai.messages": llm_gen_ai_messages,
+    # [depracated]
+    # "gen_ai.system.message": llm_gen_ai_system_message,
+    # [depracated]
+    # "gen_ai.user.message": llm_gen_ai_user_message,
+    # [depracated]
+    # "gen_ai.assistant.message": llm_gen_ai_assistant_message,
+    # -> 2.2. outputs
+    "gen_ai.choice": llm_gen_ai_choice,
+    # [debugging]
+    # "input.value": llm_input_value,
+    # [debugging]
+    # "output.value": llm_output_value,
 }

veadk/tracing/telemetry/attributes/extractors/tool_attributes_extractors.py

@@ -12,12 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
-
 from veadk.tracing.telemetry.attributes.extractors.types import (
     ExtractorResponse,
     ToolAttributesParams,
 )
+from veadk.utils.misc import safe_json_serialize
 
 
 def tool_gen_ai_operation_name(params: ToolAttributesParams) -> ExtractorResponse:
@@ -26,9 +25,8 @@ def tool_gen_ai_operation_name(params: ToolAttributesParams) -> ExtractorRespons
 
 def tool_gen_ai_tool_message(params: ToolAttributesParams) -> ExtractorResponse:
     tool_input = {
-        "id": "123",
         "role": "tool",
-        "content": json.dumps(
+        "content": safe_json_serialize(
             {
                 "name": params.tool.name,
                 "description": params.tool.description,
@@ -39,32 +37,40 @@ def tool_gen_ai_tool_message(params: ToolAttributesParams) -> ExtractorResponse:
     return ExtractorResponse(type="event", content=tool_input)
 
 
-def tool_cozeloop_input(params: ToolAttributesParams) -> ExtractorResponse:
+def tool_gen_ai_tool_input(params: ToolAttributesParams) -> ExtractorResponse:
     tool_input = {
         "name": params.tool.name,
         "description": params.tool.description,
         "parameters": params.args,
     }
-    return ExtractorResponse(content=json.dumps(tool_input) or "<unknown_tool_input>")
+    return ExtractorResponse(
+        content=safe_json_serialize(tool_input) or "<unknown_tool_input>"
+    )
 
 
 def tool_gen_ai_tool_name(params: ToolAttributesParams) -> ExtractorResponse:
     return ExtractorResponse(content=params.tool.name or "<unknown_tool_name>")
 
 
-def tool_cozeloop_output(params: ToolAttributesParams) -> ExtractorResponse:
-    function_response = params.function_response_event.get_function_responses()[0]
+def tool_gen_ai_tool_output(params: ToolAttributesParams) -> ExtractorResponse:
+    function_response = params.function_response_event.get_function_responses()[
+        0
+    ].model_dump()
     tool_output = {
-        "id": function_response.id,
-        "name": function_response.name,
-        "response": function_response.response,
+        "id": function_response["id"],
+        "name": function_response["name"],
+        "response": function_response["response"],
     }
-    return ExtractorResponse(content=json.dumps(tool_output) or "<unknown_tool_output>")
+    return ExtractorResponse(
+        content=safe_json_serialize(tool_output) or "<unknown_tool_output>"
+    )
 
 
 TOOL_ATTRIBUTES = {
     "gen_ai.operation.name": tool_gen_ai_operation_name,
     "gen_ai.tool.name": tool_gen_ai_tool_name,  # TLS required
-    "cozeloop.input": tool_cozeloop_input,  # CozeLoop required
-    "cozeloop.output": tool_cozeloop_output,  # CozeLoop required
+    "gen_ai.tool.input": tool_gen_ai_tool_input,  # TLS required
+    "gen_ai.tool.output": tool_gen_ai_tool_output,  # TLS required
+    "cozeloop.input": tool_gen_ai_tool_input,  # CozeLoop required
+    "cozeloop.output": tool_gen_ai_tool_output,  # CozeLoop required
 }

veadk/tracing/telemetry/attributes/extractors/types.py

@@ -26,13 +26,14 @@ from opentelemetry.trace.span import Span
 
 @dataclass
 class ExtractorResponse:
-    content: list | dict | None | str | int | float
+    content: Any
 
-    type: Literal["attribute", "event"] = "attribute"
+    type: Literal["attribute", "event", "event_list"] = "attribute"
     """Type of extractor response.
     
     `attribute`: span.add_attribute(attr_name, attr_value)
     `event`: span.add_event(...)
+    `event_list`: span.add_event(...) for each event in the list
     """
 
     @staticmethod
@@ -41,7 +42,7 @@ class ExtractorResponse:
     ) -> None:
         if response.type == "attribute":
             res = response.content
-            if isinstance(res, list):  # list[dict]
+            if isinstance(res, list):
                 for _res in res:
                     if isinstance(_res, dict):
                         for k, v in _res.items():
@@ -54,7 +55,17 @@ class ExtractorResponse:
                 span.add_event(attr_name, response.content)
             elif isinstance(response.content, list):
                 for event in response.content:
-                    span.add_event(attr_name, event)
+                    span.add_event(attr_name, event)  # type: ignore
+        elif response.type == "event_list":
+            if isinstance(response.content, list):
+                for event in response.content:
+                    if isinstance(event, dict):
+                        # we ensure this dict only have one key-value pair
+                        key, value = next(iter(event.items()))
+                        span.add_event(key, value)
+                    else:
+                        # Unsupported response type, discard it.
+                        pass
         else:
             # Unsupported response type, discard it.
             pass

veadk/tracing/telemetry/exporters/inmemory_exporter.py

@@ -106,3 +106,6 @@ class InMemoryExporter(BaseExporter):
 
         self._exporter = _InMemoryExporter()
         self.processor = _InMemorySpanProcessor(self._exporter)
+
+
+_INMEMORY_EXPORTER_INSTANCE = InMemoryExporter()