agentscope-runtime 0.1.5b1__py3-none-any.whl → 0.2.0b1__py3-none-any.whl

agentscope_runtime/engine/tracing/message_util.py

@@ -0,0 +1,528 @@
+# -*- coding: utf-8 -*-
+from typing import List, Optional, Union
+
+from openai.types.chat import ChatCompletionChunk
+from openai.types.chat.chat_completion_chunk import ChoiceDeltaToolCall
+
+from agentscope_runtime.engine.schemas.agent_schemas import (
+    Role,
+    FunctionCall,
+    AgentResponse,
+    RunStatus,
+    Message,
+    TextContent,
+)
+
+# Use OpenAI's ToolCall type instead of agentscope_bricks
+ToolCall = ChoiceDeltaToolCall
+
+
+# TODO: add this for streaming structured output support later
+def merge_incremental_chunk(  # pylint: disable=too-many-branches,too-many-nested-blocks  # noqa: E501
+    responses: List[ChatCompletionChunk],
+) -> Optional[ChatCompletionChunk]:
+    """
+    Merge an incremental chunk list to a ChatCompletionChunk.
+
+    Args:
+        responses (List[ChatCompletionChunk]): List of incremental chat
+            completion chunks to merge into a single response.
+
+    Returns:
+        Optional[ChatCompletionChunk]: The merged chat completion chunk,
+        or None if the input list is empty.
+    """
+
+    if len(responses) == 0:
+        return None
+
+    if not isinstance(responses[0], ChatCompletionChunk):
+        return None
+
+    # get usage or finish reason
+    merged = ChatCompletionChunk(**responses[-1].__dict__)
+
+    # if the responses has usage info, then merge the finish reason chunk to
+    # usage chunk
+    if not merged.choices and len(responses) > 1:
+        merged.choices = responses[-2].choices
+
+    # might be multiple tool calls result
+    tool_calls_dict = {}
+
+    for resp in reversed(responses[:-1]):
+        for i, j in zip(merged.choices, resp.choices):
+            # jump the finish reason chunk
+            if (i.delta.content is None and j.delta.content is not None) and (
+                i.delta.tool_calls is None and j.delta.tool_calls is not None
+            ):
+                continue
+            if j.delta.role == Role.TOOL:
+                continue
+            # merge content
+            if not i.delta.content and isinstance(j.delta.content, str):
+                i.delta.content = j.delta.content
+            elif isinstance(i.delta.content, str) and isinstance(
+                j.delta.content,
+                str,
+            ):
+                i.delta.content = j.delta.content + i.delta.content
+
+            # merge tool calls
+            elif not i.delta.tool_calls and isinstance(
+                j.delta.tool_calls,
+                list,
+            ):
+                for tool_call in j.delta.tool_calls:
+                    if tool_call.index not in tool_calls_dict:
+                        tool_calls_dict[tool_call.index] = tool_call
+                        # make sure function.arguments is a string
+                        if not tool_call.function.arguments:
+                            tool_calls_dict[
+                                tool_call.index
+                            ].function.arguments = ""
+                    else:
+                        if tool_call.id != "":
+                            tool_calls_dict[tool_call.index].id = tool_call.id
+                        if tool_call.function.name:
+                            tool_calls_dict[
+                                tool_call.index
+                            ].function.name = tool_call.function.name
+                        if (
+                            tool_call.function.arguments
+                            and not tool_calls_dict[
+                                tool_call.index
+                            ].function.arguments.startswith("{")
+                        ):
+                            tool_calls_dict[
+                                tool_call.index
+                            ].function.arguments = (
+                                tool_call.function.arguments
+                                + tool_calls_dict[
+                                    tool_call.index
+                                ].function.arguments
+                            )
+
+        if merged.usage and resp.usage:
+            merged.usage.prompt_tokens += resp.usage.prompt_tokens
+            merged.usage.completion_tokens += resp.usage.completion_tokens
+            merged.usage.total_tokens += resp.usage.total_tokens
+
+    if tool_calls_dict:
+        merged.choices[0].delta.tool_calls = [
+            ToolCall(
+                id=tool_call.id,
+                type=tool_call.type,
+                function=FunctionCall(**tool_call.function.__dict__),
+            )
+            for tool_call in tool_calls_dict.values()
+        ]
+    return merged
+
+
+def get_finish_reason(response: ChatCompletionChunk) -> Optional[str]:
+    finish_reason = None
+
+    if not isinstance(response, ChatCompletionChunk):
+        return finish_reason
+
+    if response.choices:
+        if response.choices[0].finish_reason:
+            finish_reason = response.choices[0].finish_reason
+
+    return finish_reason
+
+
+def merge_agent_response(  # pylint: disable=too-many-return-statements,too-many-branches,too-many-statements,too-many-nested-blocks  # noqa: E501
+    responses: List[Union[AgentResponse, Message, TextContent]],
+) -> AgentResponse:
+    """
+    Merge a list of incremental response objects into a single AgentResponse.
+
+    Args:
+        responses (List[Union[AgentResponse, Message, TextContent]]):
+            List of incremental responses to merge into a single response.
+
+    Returns:
+        AgentResponse: The merged agent response.
+    """
+    if len(responses) == 0:
+        raise ValueError("Cannot merge empty response list")
+
+    # Check if all responses are of the same object type
+    object_types = set()
+    for resp in responses:
+        if hasattr(resp, "object"):
+            object_types.add(resp.object)
+        else:
+            # If no object field, treat as AgentResponse
+            object_types.add("response")
+
+    if len(object_types) > 1:
+        # Mixed object types, convert the last response to AgentResponse
+        last_resp = responses[-1]
+        if isinstance(last_resp, TextContent):
+            # Convert TextContent to Message, then to AgentResponse
+            message = Message(
+                role=Role.ASSISTANT,
+                content=[last_resp],
+                status=last_resp.status or RunStatus.Completed,
+            )
+            return AgentResponse(
+                output=[message],
+                status=message.status,
+                session_id=None,
+            )
+        elif isinstance(last_resp, Message):
+            return AgentResponse(
+                output=[last_resp],
+                status=last_resp.status,
+                session_id=None,
+            )
+        else:
+            return AgentResponse(**last_resp.__dict__)
+
+    object_type = list(object_types)[0]
+
+    if object_type == "content":
+        # For content objects, merge text content and convert to AgentResponse
+        text_contents = [
+            resp for resp in responses if hasattr(resp, "text") and resp.text
+        ]
+        if not text_contents:
+            # Return empty AgentResponse if no text content
+            return AgentResponse(
+                status=RunStatus.Completed,
+                session_id=None,
+            )
+
+        # Merge all text content
+        merged_text = ""
+        last_content = text_contents[-1]
+
+        for content in text_contents:
+            if content.delta:
+                merged_text += content.text
+            else:
+                merged_text = content.text
+
+        # Create a Message with merged content
+        final_content = TextContent(
+            text=merged_text,
+            delta=False,
+            index=0,
+            msg_id=last_content.msg_id,
+            status=RunStatus.Completed,
+        )
+
+        message = Message(
+            role=Role.ASSISTANT,
+            content=[final_content],
+            status=RunStatus.Completed,
+        )
+
+        return AgentResponse(
+            output=[message],
+            status=RunStatus.Completed,
+            session_id=None,
+        )
+
+    elif object_type == "message":
+        # For message objects, convert to AgentResponse
+        messages = [resp for resp in responses if isinstance(resp, Message)]
+        if not messages:
+            return AgentResponse(
+                status=RunStatus.Completed,
+                session_id=None,
+            )
+
+        # Return the last message as AgentResponse
+        last_message = messages[-1]
+        return AgentResponse(
+            output=[last_message],
+            status=last_message.status,
+            session_id=None,
+        )
+
+    else:
+        # For response objects, use existing logic
+        # Filter only AgentResponse objects
+        agent_responses = [
+            resp for resp in responses if isinstance(resp, AgentResponse)
+        ]
+
+        if len(agent_responses) == 0:
+            last_resp = responses[-1]
+            if isinstance(last_resp, Message):
+                return AgentResponse(
+                    output=[last_resp],
+                    status=last_resp.status,
+                    session_id=None,
+                )
+            else:
+                return AgentResponse(**last_resp.__dict__)
+
+        # Get the last AgentResponse as base
+        merged = AgentResponse(**agent_responses[-1].__dict__)
+
+        # If no output, return the merged response
+        if not merged.output:
+            return merged
+
+        # Merge content from all AgentResponse objects
+        content_dict = {}
+
+        for resp in agent_responses:
+            if not resp.output:
+                continue
+
+            for message in resp.output:
+                if not message.content:
+                    continue
+
+                for content in message.content:
+                    if (
+                        content.type == "text"
+                        and hasattr(content, "text")
+                        and content.text
+                    ):
+                        # For text content, accumulate the text
+                        if content.msg_id not in content_dict:
+                            content_dict[content.msg_id] = {
+                                "content": content,
+                                "text": content.text,
+                                "delta": content.delta,
+                            }
+                        else:
+                            # If delta is True, append text; if False, replace
+                            if content.delta:
+                                content_dict[content.msg_id][
+                                    "text"
+                                ] += content.text
+                            else:
+                                content_dict[content.msg_id][
+                                    "text"
+                                ] = content.text
+
+                            # Update the content object with merged text
+                            content_dict[content.msg_id][
+                                "content"
+                            ].text = content_dict[content.msg_id]["text"]
+                            content_dict[content.msg_id][
+                                "content"
+                            ].delta = False
+
+        # Update the merged response with accumulated content
+        if content_dict:
+            for message in merged.output:
+                if message.content:
+                    for content in message.content:
+                        if (
+                            content.type == "text"
+                            and hasattr(content, "msg_id")
+                            and content.msg_id in content_dict
+                        ):
+                            content.text = content_dict[content.msg_id]["text"]
+                            content.delta = False
+
+        return merged
+
+
+def get_agent_response_finish_reason(
+    response: Union[AgentResponse, Message, TextContent],
+) -> Optional[str]:
+    """
+    Get the finish reason from a response object.
+
+    Args:
+        response (Union[AgentResponse, Message, TextContent]):
+            The response object
+
+    Returns:
+        Optional[str]: The finish reason, or None if not finished
+    """
+    # Only consider AgentResponse objects as potential final responses
+    if isinstance(response, AgentResponse):
+        if (
+            hasattr(response, "status")
+            and response.status == RunStatus.Completed
+        ):
+            # Check if this is a final response with output
+            if hasattr(response, "output") and response.output:
+                return "stop"
+    return None
+
+
+def merge_agent_message(  # pylint: disable=too-many-return-statements,too-many-branches,too-many-statements,too-many-nested-blocks  # noqa: E501
+    messages: List[Union[Message, TextContent]],
+) -> Message:
+    """
+    Merge a list of incremental message objects into a single Message.
+
+    Args:
+        messages (List[Union[Message, TextContent]]):
+            List of incremental messages to merge into a single message.
+
+    Returns:
+        Message: The merged message.
+    """
+    if len(messages) == 0:
+        raise ValueError("Cannot merge empty message list")
+
+    # Check if all messages are of the same object type
+    object_types = set()
+    for msg in messages:
+        if hasattr(msg, "object"):
+            object_types.add(msg.object)
+        else:
+            # If no object field, treat as Message
+            object_types.add("message")
+
+    if len(object_types) > 1:
+        # Mixed object types, convert the last message to Message
+        last_msg = messages[-1]
+        if isinstance(last_msg, TextContent):
+            # Convert TextContent to Message with delta=False
+            final_content = TextContent(
+                text=last_msg.text,
+                delta=False,
+                index=last_msg.index,
+                msg_id=last_msg.msg_id,
+                status=RunStatus.Completed,
+            )
+            return Message(
+                role=Role.ASSISTANT,
+                content=[final_content],
+                status=RunStatus.Completed,
+            )
+        else:
+            return Message(**last_msg.__dict__)
+
+    object_type = list(object_types)[0]
+
+    if object_type == "content":
+        # For content objects, merge text content and convert to Message
+        text_contents = [
+            msg for msg in messages if hasattr(msg, "text") and msg.text
+        ]
+        if not text_contents:
+            # Return empty Message if no text content
+            return Message(
+                role=Role.ASSISTANT,
+                status=RunStatus.Completed,
+            )
+
+        # Merge all text content
+        merged_text = ""
+        last_content = text_contents[-1]
+
+        for content in text_contents:
+            if content.delta:
+                merged_text += content.text
+            else:
+                merged_text = content.text
+
+        # Create a Message with merged content
+        final_content = TextContent(
+            text=merged_text,
+            delta=False,
+            index=0,
+            msg_id=last_content.msg_id,
+            status=RunStatus.Completed,
+        )
+
+        return Message(
+            role=Role.ASSISTANT,
+            content=[final_content],
+            status=RunStatus.Completed,
+        )
+
+    else:
+        # For message objects, use existing logic
+        # Filter only Message objects
+        message_objects = [msg for msg in messages if isinstance(msg, Message)]
+
+        if len(message_objects) == 0:
+            last_msg = messages[-1]
+            if isinstance(last_msg, TextContent):
+                return Message(
+                    role=Role.ASSISTANT,
+                    content=[last_msg],
+                    status=last_msg.status or RunStatus.Completed,
+                )
+            else:
+                return Message(**last_msg.__dict__)
+
+        # Get the last Message as base
+        merged = Message(**message_objects[-1].__dict__)
+
+        # If no content, return the merged message
+        if not merged.content:
+            return merged
+
+        # Merge content from all Message objects
+        # Use msg_id + index as key to avoid overwriting different contents
+        # with the same msg_id but different indexes
+        content_dict = {}
+
+        for msg in message_objects:
+            if not msg.content:
+                continue
+
+            for content in msg.content:
+                if (
+                    content.type == "text"
+                    and hasattr(content, "text")
+                    and content.text
+                ):
+                    # Create unique key using msg_id and index
+                    content_key = f"{content.msg_id}_{content.index}"
+
+                    # For text content, accumulate the text
+                    if content_key not in content_dict:
+                        content_dict[content_key] = {
+                            "content": content,
+                            "text": content.text,
+                            "delta": content.delta,
+                            "msg_id": content.msg_id,
+                            "index": content.index,
+                        }
+                    else:
+                        # If delta is True, append text; if False, replace
+                        if content.delta:
+                            content_dict[content_key]["text"] += content.text
+                        else:
+                            content_dict[content_key]["text"] = content.text
+
+                        # Update the content object with merged text
+                        content_dict[content_key][
+                            "content"
+                        ].text = content_dict[content_key]["text"]
+                        content_dict[content_key]["content"].delta = False
+
+        # Update the merged message with accumulated content
+        if content_dict:
+            for content in merged.content:
+                if (
+                    content.type == "text"
+                    and hasattr(content, "msg_id")
+                    and hasattr(content, "index")
+                ):
+                    content_key = f"{content.msg_id}_{content.index}"
+                    if content_key in content_dict:
+                        content.text = content_dict[content_key]["text"]
+                        content.delta = False
+
+        return merged
+
+
+def get_agent_message_finish_reason(
+    message: Optional[Union[Message, TextContent]],
+) -> Optional[str]:
+    if message is None:
+        return None
+
+    if isinstance(message, Message):
+        return "stop" if message.status == RunStatus.Completed else None
+
+    return None

agentscope_runtime/engine/tracing/tracing_metric.py

@@ -10,14 +10,26 @@ class TraceType(str):
     """
 
     # Officially supported event types
-    LLM = "llm"
-    TOOL = "tool"
-    AGENT_STEP = "agent_step"
-    SEARCH = "search"
-    IMAGE_GENERATION = "image_generation"
-    RAG = "rag"
-    INTENTION = "intention"
-    PLUGIN_CENTER = "plugin_center"
+    LLM = "LLM"
+    TOOL = "TOOL"
+    AGENT_STEP = "AGENT_STEP"
+    MEMORY = "MEMORY"
+    SEARCH = "SEARCH"
+    AIGC = "AIGC"
+    RAG = "RAG"
+    INTENTION = "INTENTION"
+    PLUGIN = "PLUGIN"
+    FUNCTION_CALL = "FUNCTION_CALL"
+    AGENT = "AGENT"
+    PLANNING = "PLANNING"
+    CHAIN = "CHAIN"
+    RETRIEVER = "RETRIEVER"
+    RERANKER = "RERANKER"
+    EMBEDDING = "EMBEDDING"
+    TASK = "TASK"
+    GUARDRAIL = "GUARDRAIL"
+    REWRITER = "REWRITER"
+    OTHER = "OTHER"
 
     def __init__(self, value: str):
         """Initialize a TraceType with a string value.

agentscope_runtime/engine/tracing/tracing_util.py

@@ -0,0 +1,130 @@
+# -*- coding: utf-8 -*-
+import contextvars
+import os
+from opentelemetry import baggage
+from opentelemetry.context import attach
+
+_user_request_id: contextvars.ContextVar[str] = contextvars.ContextVar(
+    "_user_request_id",
+    default="",
+)
+
+_user_trace_header: contextvars.ContextVar[dict] = contextvars.ContextVar(
+    "_user_trace_header",
+    default={},
+)
+
+_user_common_attributes: contextvars.ContextVar[dict] = contextvars.ContextVar(
+    "_user_common_attributes",
+    default={},
+)
+
+
+class TracingUtil:
+    @staticmethod
+    def set_request_id(value: str = "") -> None:
+        """Set request id"""
+        _user_request_id.set(value)
+        if value:
+            ctx = baggage.set_baggage("agentscope-bricks_request_id", value)
+            attach(ctx)
+
+            TracingUtil.clear_common_attributes()
+            TracingUtil.set_common_attributes(
+                {
+                    "request_id": value,
+                    "bailian.request_id": value,
+                    "gen_ai.response.id": value,
+                    **_global_attributes,
+                },
+            )
+
+    @staticmethod
+    def get_request_id(default_value: str = "") -> str:
+        """Get request id"""
+        # Try to get from context variable first
+        request_id = _user_request_id.get("")
+        if request_id:
+            return request_id
+
+        # Fallback to baggage for cross-thread scenarios
+        request_id = baggage.get_baggage("agentscope-bricks_request_id")
+        if request_id:
+            return request_id
+
+        return default_value
+
+    @staticmethod
+    def set_trace_header(trace_headers: dict) -> None:
+        """Set trace headers
+
+        Args:
+            trace_headers: trace headers to set
+        """
+        _user_trace_header.set(trace_headers)
+
+    @staticmethod
+    def get_trace_header() -> dict:
+        """Get trace headers
+
+        Returns:
+            trace headers in dict
+        """
+        return _user_trace_header.get({})
+
+    @staticmethod
+    def set_common_attributes(attributes: dict) -> None:
+        """Set common attributes by merging with existing ones
+
+        Args:
+            attributes: common attributes to merge
+        """
+        current_attributes = _user_common_attributes.get({})
+        current_attributes.update(attributes)
+        _user_common_attributes.set(current_attributes)
+
+    @staticmethod
+    def get_common_attributes() -> dict:
+        """Get common attributes
+
+        Returns:
+            common attributes in dict
+        """
+        return _user_common_attributes.get({})
+
+    @staticmethod
+    def clear_common_attributes() -> None:
+        """Clear common attributes"""
+        _user_common_attributes.set({})
+
+
+def get_global_attributes() -> dict:
+    """Set global common attributes for tracing."""
+    attributes = {"gen_ai.framework": "Alibaba Cloud Model Studio"}
+
+    if app_env := (os.getenv("APPLICATION_ENV") or os.getenv("DS_ENV")):
+        attributes["bailian.app.env"] = app_env
+
+    if app_id := os.getenv("APPLICATION_ID"):
+        attributes["bailian.app.id"] = app_id
+
+    if app_name := os.getenv("APPLICATION_NAME"):
+        attributes["bailian.app.name"] = app_name
+
+    if app_inter_source := os.getenv("APPLICATION_INTER_SOURCE"):
+        attributes["bailian.app.inter.source"] = app_inter_source
+
+    if user_id := os.getenv("ALIYUN_UID"):
+        attributes["bailian.app.owner_id"] = user_id
+        attributes["gen_ai.user.id"] = user_id
+
+    if app_tracing := os.getenv("APPLICATION_TRACING"):
+        attributes["bailian.app.tracing"] = app_tracing
+
+    if workspace_id := os.getenv("WORKSPACE_ID"):
+        attributes["bailian.app.workspace"] = workspace_id
+
+    return attributes
+
+
+_global_attributes = get_global_attributes()