veadk-python 0.2.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

veadk/runner.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import functools
+import os
 from types import MethodType
 from typing import Union
 
@@ -34,7 +34,7 @@ from veadk.evaluation import EvalSetRecorder
 from veadk.memory.short_term_memory import ShortTermMemory
 from veadk.types import MediaMessage
 from veadk.utils.logger import get_logger
-from veadk.utils.misc import formatted_timestamp, read_png_to_bytes
+from veadk.utils.misc import formatted_timestamp, read_file_to_bytes
 
 logger = get_logger(__name__)
 
@@ -50,11 +50,7 @@ RunnerMessage = Union[
 async def pre_run_process(self, process_func, new_message, user_id, session_id):
     if new_message.parts:
         for part in new_message.parts:
-            if (
-                part.inline_data
-                and part.inline_data.mime_type == "image/png"
-                and self.upload_inline_data_to_tos
-            ):
+            if part.inline_data and self.upload_inline_data_to_tos:
                 await process_func(
                     part,
                     self.app_name,
@@ -105,9 +101,20 @@ def _convert_messages(
     if isinstance(messages, str):
         _messages = [types.Content(role="user", parts=[types.Part(text=messages)])]
     elif isinstance(messages, MediaMessage):
-        assert messages.media.endswith(".png"), (
-            "The MediaMessage only supports PNG format file for now."
+        import filetype
+
+        file_data = read_file_to_bytes(messages.media)
+
+        kind = filetype.guess(file_data)
+        if kind is None:
+            raise ValueError("Unsupported or unknown file type.")
+
+        mime_type = kind.mime
+
+        assert mime_type.startswith(("image/", "video/")), (
+            f"Unsupported media type: {mime_type}"
         )
+
         _messages = [
             types.Content(
                 role="user",
@@ -116,8 +123,8 @@ def _convert_messages(
                     types.Part(
                         inline_data=Blob(
                             display_name=messages.media,
-                            data=read_png_to_bytes(messages.media),
-                            mime_type="image/png",
+                            data=file_data,
+                            mime_type=mime_type,
                         )
                     ),
                 ],
@@ -249,15 +256,21 @@ class Runner(ADKRunner):
         )
 
         if self.short_term_memory:
-            await self.short_term_memory.create_session(
-                app_name=self.app_name, user_id=self.user_id, session_id=session_id
+            session = await self.short_term_memory.create_session(
+                app_name=self.app_name, user_id=user_id, session_id=session_id
+            )
+            assert session, (
+                f"Failed to create session with app_name={self.app_name}, user_id={user_id}, session_id={session_id}, "
+            )
+            logger.debug(
+                f"Auto create session: {session.id}, user_id: {session.user_id}, app_name: {self.app_name}"
             )
 
         final_output = ""
         for converted_message in converted_messages:
             try:
                 async for event in self.run_async(
-                    user_id=self.user_id,
+                    user_id=user_id,
                     session_id=session_id,
                     new_message=converted_message,
                     run_config=run_config,
@@ -271,7 +284,8 @@ class Runner(ADKRunner):
                         and event.content.parts[0].text is not None
                         and len(event.content.parts[0].text.strip()) > 0
                     ):
-                        final_output += event.content.parts[0].text
+                        final_output = event.content.parts[0].text
+                        logger.debug(f"Event output: {final_output}")
             except LlmCallsLimitExceededError as e:
                 logger.warning(f"Max number of llm calls limit exceeded: {e}")
                 final_output = ""
@@ -359,19 +373,28 @@ class Runner(ADKRunner):
         )
         return eval_set_path
 
-    async def save_session_to_long_term_memory(self, session_id: str) -> None:
+    async def save_session_to_long_term_memory(
+        self, session_id: str, user_id: str = "", app_name: str = ""
+    ) -> None:
         if not self.long_term_memory:
             logger.warning("Long-term memory is not enabled. Failed to save session.")
             return
 
+        if not user_id:
+            user_id = self.user_id
+
+        if not app_name:
+            app_name = self.app_name
+
         session = await self.session_service.get_session(
-            app_name=self.app_name,
-            user_id=self.user_id,
+            app_name=app_name,
+            user_id=user_id,
             session_id=session_id,
         )
+
         if not session:
             logger.error(
-                f"Session {session_id} not found in session service, cannot save to long-term memory."
+                f"Session {session_id} (app_name={app_name}, user_id={user_id}) not found in session service, cannot save to long-term memory."
             )
             return
 

veadk/tools/builtin_tools/generate_image.py

@@ -12,28 +12,35 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import base64
+import json
+import mimetypes
+import traceback
 from typing import Dict
 
 from google.adk.tools import ToolContext
-from veadk.config import getenv
-from veadk.consts import DEFAULT_IMAGE_GENERATE_MODEL_NAME, DEFAULT_MODEL_AGENT_API_BASE
-
-import base64
-from volcenginesdkarkruntime import Ark
+from google.genai.types import Blob, Part
 from opentelemetry import trace
-import traceback
-from veadk.version import VERSION
 from opentelemetry.trace import Span
-from veadk.utils.logger import get_logger
+from volcenginesdkarkruntime import Ark
 from volcenginesdkarkruntime.types.images.images import SequentialImageGenerationOptions
-import json
 
+from veadk.config import getenv, settings
+from veadk.consts import (
+    DEFAULT_IMAGE_GENERATE_MODEL_NAME,
+    DEFAULT_IMAGE_GENERATE_MODEL_API_BASE,
+)
+from veadk.utils.logger import get_logger
+from veadk.utils.misc import formatted_timestamp, read_file_to_bytes
+from veadk.version import VERSION
 
 logger = get_logger(__name__)
 
 client = Ark(
-    api_key=getenv("MODEL_AGENT_API_KEY"),
-    base_url=DEFAULT_MODEL_AGENT_API_BASE,
+    api_key=getenv(
+        "MODEL_IMAGE_API_KEY", getenv("MODEL_AGENT_API_KEY", settings.model.api_key)
+    ),
+    base_url=getenv("MODEL_IMAGE_API_BASE", DEFAULT_IMAGE_GENERATE_MODEL_API_BASE),
 )
 
 
@@ -121,7 +128,7 @@ async def image_generate(
     - size 推荐使用 2048x2048 或表格里的标准比例，确保生成质量。
     """
 
-    success_list = []
+    success_list: list[dict] = []
     error_list = []
 
     for idx, item in enumerate(tasks):
@@ -184,7 +191,9 @@ async def image_generate(
                     and max_images
                 ):
                     response = client.images.generate(
-                        model=DEFAULT_IMAGE_GENERATE_MODEL_NAME,
+                        model=getenv(
+                            "MODEL_IMAGE_NAME", DEFAULT_IMAGE_GENERATE_MODEL_NAME
+                        ),
                         **inputs,
                         sequential_image_generation_options=SequentialImageGenerationOptions(
                             max_images=max_images
@@ -192,7 +201,10 @@ async def image_generate(
                     )
                 else:
                     response = client.images.generate(
-                        model=DEFAULT_IMAGE_GENERATE_MODEL_NAME, **inputs
+                        model=getenv(
+                            "MODEL_IMAGE_NAME", DEFAULT_IMAGE_GENERATE_MODEL_NAME
+                        ),
+                        **inputs,
                     )
                 if not response.error:
                     for i, image_data in enumerate(response.data):
@@ -261,8 +273,12 @@ async def image_generate(
                 output_part=output_part,
                 output_tokens=output_tokens,
                 total_tokens=total_tokens,
-                request_model=DEFAULT_IMAGE_GENERATE_MODEL_NAME,
-                response_model=DEFAULT_IMAGE_GENERATE_MODEL_NAME,
+                request_model=getenv(
+                    "MODEL_IMAGE_NAME", DEFAULT_IMAGE_GENERATE_MODEL_NAME
+                ),
+                response_model=getenv(
+                    "MODEL_IMAGE_NAME", DEFAULT_IMAGE_GENERATE_MODEL_NAME
+                ),
             )
     if len(success_list) == 0:
         return {
@@ -271,6 +287,28 @@ async def image_generate(
             "error_list": error_list,
         }
     else:
+        app_name = tool_context._invocation_context.app_name
+        user_id = tool_context._invocation_context.user_id
+        session_id = tool_context._invocation_context.session.id
+
+        artifact_service = tool_context._invocation_context.artifact_service
+        if artifact_service:
+            for image in success_list:
+                for _, image_tos_url in image.items():
+                    filename = f"artifact_{formatted_timestamp()}"
+                    await artifact_service.save_artifact(
+                        app_name=app_name,
+                        user_id=user_id,
+                        session_id=session_id,
+                        filename=filename,
+                        artifact=Part(
+                            inline_data=Blob(
+                                display_name=filename,
+                                data=read_file_to_bytes(image_tos_url),
+                                mime_type=mimetypes.guess_type(image_tos_url)[0],
+                            )
+                        ),
+                    )
         return {
             "status": "success",
             "success_list": success_list,
@@ -332,10 +370,11 @@ def add_span_attributes(
 
 def _upload_image_to_tos(image_bytes: bytes, object_key: str) -> None:
     try:
-        from veadk.integrations.ve_tos.ve_tos import VeTOS
         import os
         from datetime import datetime
 
+        from veadk.integrations.ve_tos.ve_tos import VeTOS
+
         timestamp: str = datetime.now().strftime("%Y%m%d%H%M%S%f")[:-3]
         object_key = f"{timestamp}-{object_key}"
         bucket_name = os.getenv("DATABASE_TOS_BUCKET")

veadk/tools/builtin_tools/image_edit.py

@@ -15,8 +15,11 @@
 from typing import Dict
 from google.adk.tools import ToolContext
 from volcenginesdkarkruntime import Ark
-from veadk.config import getenv
-from veadk.consts import DEFAULT_MODEL_AGENT_API_BASE, DEFAULT_IMAGE_EDIT_MODEL_NAME
+from veadk.config import getenv, settings
+from veadk.consts import (
+    DEFAULT_IMAGE_EDIT_MODEL_API_BASE,
+    DEFAULT_IMAGE_EDIT_MODEL_NAME,
+)
 import base64
 from opentelemetry import trace
 import traceback
@@ -28,8 +31,10 @@ from veadk.utils.logger import get_logger
 logger = get_logger(__name__)
 
 client = Ark(
-    api_key=getenv("MODEL_AGENT_API_KEY"),
-    base_url=DEFAULT_MODEL_AGENT_API_BASE,
+    api_key=getenv(
+        "MODEL_EDIT_API_KEY", getenv("MODEL_AGENT_API_KEY", settings.model.api_key)
+    ),
+    base_url=getenv("MODEL_EDIT_API_BASE", DEFAULT_IMAGE_EDIT_MODEL_API_BASE),
 )
 
 
@@ -123,7 +128,8 @@ async def image_edit(
                     "parts.1.image_url.url": origin_image,
                 }
                 response = client.images.generate(
-                    model=DEFAULT_IMAGE_EDIT_MODEL_NAME, **inputs
+                    model=getenv("MODEL_EDIT_NAME", DEFAULT_IMAGE_EDIT_MODEL_NAME),
+                    **inputs,
                 )
                 output_part = None
                 if response.data and len(response.data) > 0:
@@ -175,8 +181,12 @@ async def image_edit(
                     output_part=output_part,
                     output_tokens=response.usage.output_tokens,
                     total_tokens=response.usage.total_tokens,
-                    request_model=DEFAULT_IMAGE_EDIT_MODEL_NAME,
-                    response_model=DEFAULT_IMAGE_EDIT_MODEL_NAME,
+                    request_model=getenv(
+                        "MODEL_EDIT_NAME", DEFAULT_IMAGE_EDIT_MODEL_NAME
+                    ),
+                    response_model=getenv(
+                        "MODEL_EDIT_NAME", DEFAULT_IMAGE_EDIT_MODEL_NAME
+                    ),
                 )
 
         except Exception as e:

veadk/tools/builtin_tools/image_generate.py

@@ -15,8 +15,11 @@
 from typing import Dict
 
 from google.adk.tools import ToolContext
-from veadk.config import getenv
-from veadk.consts import DEFAULT_TEXT_TO_IMAGE_MODEL_NAME, DEFAULT_MODEL_AGENT_API_BASE
+from veadk.config import getenv, settings
+from veadk.consts import (
+    DEFAULT_TEXT_TO_IMAGE_MODEL_NAME,
+    DEFAULT_TEXT_TO_IMAGE_MODEL_API_BASE,
+)
 import base64
 from volcenginesdkarkruntime import Ark
 from opentelemetry import trace
@@ -29,8 +32,10 @@ from veadk.utils.logger import get_logger
 logger = get_logger(__name__)
 
 client = Ark(
-    api_key=getenv("MODEL_AGENT_API_KEY"),
-    base_url=DEFAULT_MODEL_AGENT_API_BASE,
+    api_key=getenv(
+        "MODEL_IMAGE_API_KEY", getenv("MODEL_AGENT_API_KEY", settings.model.api_key)
+    ),
+    base_url=getenv("MODEL_IMAGE_API_BASE", DEFAULT_TEXT_TO_IMAGE_MODEL_API_BASE),
 )
 
 
@@ -120,7 +125,8 @@ async def image_generate(
                     "content": json.dumps(inputs, ensure_ascii=False),
                 }
                 response = client.images.generate(
-                    model=DEFAULT_TEXT_TO_IMAGE_MODEL_NAME, **inputs
+                    model=getenv("MODEL_IMAGE_NAME", DEFAULT_TEXT_TO_IMAGE_MODEL_NAME),
+                    **inputs,
                 )
                 output_part = None
                 if response.data and len(response.data) > 0:
@@ -172,8 +178,12 @@ async def image_generate(
                     output_part=output_part,
                     output_tokens=response.usage.output_tokens,
                     total_tokens=response.usage.total_tokens,
-                    request_model=DEFAULT_TEXT_TO_IMAGE_MODEL_NAME,
-                    response_model=DEFAULT_TEXT_TO_IMAGE_MODEL_NAME,
+                    request_model=getenv(
+                        "MODEL_IMAGE_NAME", DEFAULT_TEXT_TO_IMAGE_MODEL_NAME
+                    ),
+                    response_model=getenv(
+                        "MODEL_IMAGE_NAME", DEFAULT_TEXT_TO_IMAGE_MODEL_NAME
+                    ),
                 )
 
         except Exception as e:

veadk/tools/builtin_tools/load_knowledgebase.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from google.adk.models.llm_request import LlmRequest
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+from pydantic import BaseModel, Field
+from typing_extensions import override
+
+from veadk.knowledgebase import KnowledgeBase
+from veadk.knowledgebase.entry import KnowledgebaseEntry
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class LoadKnowledgebaseResponse(BaseModel):
+    knowledges: list[KnowledgebaseEntry] = Field(default_factory=list)
+
+
+class LoadKnowledgebaseTool(FunctionTool):
+    """A tool that loads the common knowledgebase"""
+
+    def __init__(self, knowledgebase: KnowledgeBase):
+        super().__init__(self.load_knowledgebase)
+
+        self.knowledgebase = knowledgebase
+
+        if not self.custom_metadata:
+            self.custom_metadata = {}
+        self.custom_metadata["backend"] = knowledgebase.backend
+
+    @override
+    def _get_declaration(self) -> types.FunctionDeclaration | None:
+        return types.FunctionDeclaration(
+            name=self.name,
+            description=self.description,
+            parameters=types.Schema(
+                type=types.Type.OBJECT,
+                properties={
+                    "query": types.Schema(
+                        type=types.Type.STRING,
+                    )
+                },
+                required=["query"],
+            ),
+        )
+
+    @override
+    async def process_llm_request(
+        self,
+        *,
+        tool_context: ToolContext,
+        llm_request: LlmRequest,
+    ) -> None:
+        await super().process_llm_request(
+            tool_context=tool_context, llm_request=llm_request
+        )
+        # Tell the model about the knowledgebase.
+        llm_request.append_instructions(
+            [
+                f"""
+You have a knowledgebase (knowledegebase name is `{self.knowledgebase.name}`, knowledgebase description is `{self.knowledgebase.description}`). You can use it to answer questions. If any questions need
+you to look up the knowledgebase, you should call load_knowledgebase function with a query.
+"""
+            ]
+        )
+
+    async def load_knowledgebase(
+        self, query: str, tool_context: ToolContext
+    ) -> LoadKnowledgebaseResponse:
+        """Loads the knowledgebase for the user.
+
+        Args:
+        query: The query to load the knowledgebase for.
+
+        Returns:
+        A list of knowledgebase results.
+        """
+        logger.info(f"Search knowledgebase: {self.knowledgebase.name}")
+        response = self.knowledgebase.search(query)
+        logger.info(f"Loaded {len(response)} knowledgebase entries for query: {query}")
+        return LoadKnowledgebaseResponse(knowledges=response)

veadk/tools/builtin_tools/video_generate.py

@@ -25,16 +25,18 @@ from volcenginesdkarkruntime.types.content_generation.create_task_content_param
     CreateTaskContentParam,
 )
 
-from veadk.config import getenv
-from veadk.consts import DEFAULT_MODEL_AGENT_API_BASE, DEFAULT_VIDEO_MODEL_NAME
+from veadk.config import getenv, settings
+from veadk.consts import DEFAULT_VIDEO_MODEL_API_BASE, DEFAULT_VIDEO_MODEL_NAME
 from veadk.utils.logger import get_logger
 from veadk.version import VERSION
 
 logger = get_logger(__name__)
 
 client = Ark(
-    api_key=getenv("MODEL_AGENT_API_KEY"),
-    base_url=DEFAULT_MODEL_AGENT_API_BASE,
+    api_key=getenv(
+        "MODEL_VIDEO_API_KEY", getenv("MODEL_AGENT_API_KEY", settings.model.api_key)
+    ),
+    base_url=getenv("MODEL_VIDEO_API_BASE", DEFAULT_VIDEO_MODEL_API_BASE),
 )
 
 
@@ -43,7 +45,7 @@ async def generate(prompt, first_frame_image=None, last_frame_image=None):
         if first_frame_image is None:
             logger.debug("text generation")
             response = client.content_generation.tasks.create(
-                model=DEFAULT_VIDEO_MODEL_NAME,
+                model=getenv("MODEL_VIDEO_NAME", DEFAULT_VIDEO_MODEL_NAME),
                 content=[
                     {"type": "text", "text": prompt},
                 ],
@@ -51,7 +53,7 @@ async def generate(prompt, first_frame_image=None, last_frame_image=None):
         elif last_frame_image is None:
             logger.debug("first frame generation")
             response = client.content_generation.tasks.create(
-                model=DEFAULT_VIDEO_MODEL_NAME,
+                model=getenv("MODEL_VIDEO_NAME", DEFAULT_VIDEO_MODEL_NAME),
                 content=cast(
                     list[CreateTaskContentParam],  # avoid IDE warning
                     [
@@ -66,7 +68,7 @@ async def generate(prompt, first_frame_image=None, last_frame_image=None):
         else:
             logger.debug("last frame generation")
             response = client.content_generation.tasks.create(
-                model=DEFAULT_VIDEO_MODEL_NAME,
+                model=getenv("MODEL_VIDEO_NAME", DEFAULT_VIDEO_MODEL_NAME),
                 content=[
                     {"type": "text", "text": prompt},
                     {
@@ -263,8 +265,8 @@ async def video_generate(params: list, tool_context: ToolContext) -> Dict:
                 output_part=output_part,
                 output_tokens=total_tokens,
                 total_tokens=total_tokens,
-                request_model=DEFAULT_VIDEO_MODEL_NAME,
-                response_model=DEFAULT_VIDEO_MODEL_NAME,
+                request_model=getenv("MODEL_VIDEO_NAME", DEFAULT_VIDEO_MODEL_NAME),
+                response_model=getenv("MODEL_VIDEO_NAME", DEFAULT_VIDEO_MODEL_NAME),
             )
 
     if len(success_list) == 0:

veadk/tools/builtin_tools/web_search.py

@@ -23,6 +23,7 @@ import json
 from urllib.parse import quote
 
 import requests
+from google.adk.tools import ToolContext
 
 from veadk.config import getenv
 from veadk.utils.logger import get_logger
@@ -151,7 +152,7 @@ def request(method, date, query, header, ak, sk, action, body):
     return r.json()
 
 
-def web_search(query: str) -> list[str]:
+def web_search(query: str, tool_context: ToolContext) -> list[str]:
     """Search a query in websites.
 
     Args:
@@ -166,8 +167,14 @@ def web_search(query: str) -> list[str]:
         "Count": 5,
         "NeedSummary": True,
     }
-    ak = getenv("VOLCENGINE_ACCESS_KEY")
-    sk = getenv("VOLCENGINE_SECRET_KEY")
+
+    ak = tool_context.state.get("VOLCENGINE_ACCESS_KEY")
+    if not ak:
+        ak = getenv("VOLCENGINE_ACCESS_KEY")
+
+    sk = tool_context.state.get("VOLCENGINE_SECRET_KEY")
+    if not sk:
+        sk = getenv("VOLCENGINE_SECRET_KEY")
 
     now = datetime.datetime.utcnow()
     response_body = request(

veadk/tools/load_knowledgebase_tool.py

@@ -25,6 +25,9 @@ from typing_extensions import override
 
 from veadk.knowledgebase import KnowledgeBase
 from veadk.knowledgebase.entry import KnowledgebaseEntry
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
 
 if TYPE_CHECKING:
     from google.adk.models.llm_request import LlmRequest
@@ -96,6 +99,15 @@ class LoadKnowledgebaseTool(FunctionTool):
 
     def __init__(self):
         super().__init__(load_knowledgebase)
+        global knowledgebase
+        if knowledgebase is None:
+            logger.info(
+                "Get global knowledgebase instance failed, failed to set knowledgebase tool backend."
+            )
+        else:
+            if not self.custom_metadata:
+                self.custom_metadata = {}
+            self.custom_metadata["backend"] = knowledgebase.backend
 
     @override
     def _get_declaration(self) -> types.FunctionDeclaration | None:

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

@@ -325,6 +325,10 @@ def llm_gen_ai_operation_name(params: LLMAttributesParams) -> ExtractorResponse:
     return ExtractorResponse(content="chat")
 
 
+def llm_gen_ai_span_kind(params: LLMAttributesParams) -> ExtractorResponse:
+    return ExtractorResponse(content="llm")
+
+
 # def llm_gen_ai_system_message(params: LLMAttributesParams) -> ExtractorResponse:
 #     event_attributes = {
 #         "content": str(params.llm_request.config.system_instruction),
@@ -559,6 +563,7 @@ LLM_ATTRIBUTES = {
     "gen_ai.is_streaming": llm_gen_ai_is_streaming,
     # -> 1.4. span kind
     "gen_ai.operation.name": llm_gen_ai_operation_name,
+    "gen_ai.span.kind": llm_gen_ai_span_kind,  # apmplus required
     # -> 1.5. inputs
     "gen_ai.prompt": llm_gen_ai_prompt,
     # -> 1.6. outputs

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

@@ -23,6 +23,10 @@ def tool_gen_ai_operation_name(params: ToolAttributesParams) -> ExtractorRespons
     return ExtractorResponse(content="execute_tool")
 
 
+def tool_gen_ai_span_kind(params: ToolAttributesParams) -> ExtractorResponse:
+    return ExtractorResponse(content="tool")
+
+
 def tool_gen_ai_tool_message(params: ToolAttributesParams) -> ExtractorResponse:
     tool_input = {
         "role": "tool",
@@ -73,4 +77,7 @@ TOOL_ATTRIBUTES = {
     "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
+    "gen_ai.span.kind": tool_gen_ai_span_kind,  # apmplus required
+    "gen_ai.input": tool_gen_ai_tool_input,  # apmplus required
+    "gen_ai.output": tool_gen_ai_tool_output,  # apmplus required
 }