Merge branch 'main' of https://github.com/doraemonlove/veadk-python into fix/builtintooltrace

Author: wangjiaju.716

config.yaml.full

@@ -54,6 +54,15 @@ tool:
     dataset_id:  #dataset name
   mcp_router:
     url: #mcp sse/streamable-http url
+  code_sandbox:
+    url: #mcp sse/streamable-http url
+    api_key: #mcp api key
+  browser_sandbox:
+    url: #mcp sse/streamable-http url
+    api_key: #mcp api key
+  computer_sandbox:
+    url: #mcp sse/streamable-http url
+    api_key: #mcp api key
 
 
 observability:

veadk/runner.py

@@ -14,6 +14,7 @@
 from typing import Union
 
 from google.adk.agents import RunConfig
+from google.adk.agents.invocation_context import LlmCallsLimitExceededError
 from google.adk.agents.run_config import StreamingMode
 from google.adk.plugins.base_plugin import BasePlugin
 from google.adk.runners import Runner as ADKRunner
@@ -49,20 +50,25 @@ class Runner:
     def __init__(
         self,
         agent: VeAgent,
-        short_term_memory: ShortTermMemory,
+        short_term_memory: ShortTermMemory | None = None,
         plugins: list[BasePlugin] | None = None,
         app_name: str = "veadk_default_app",
         user_id: str = "veadk_default_user",
     ):
-        # basic settings
         self.app_name = app_name
         self.user_id = user_id
 
-        # agent settings
         self.agent = agent
 
-        self.short_term_memory = short_term_memory
-        self.session_service = short_term_memory.session_service
+        if not short_term_memory:
+            logger.info(
+                "No short term memory provided, using a in-memory memory by default."
+            )
+            self.short_term_memory = ShortTermMemory()
+        else:
+            self.short_term_memory = short_term_memory
+
+        self.session_service = self.short_term_memory.session_service
 
         # prevent VeRemoteAgent has no long-term memory attr
         if isinstance(self.agent, Agent):
@@ -114,35 +120,44 @@ async def _run(
         self,
         session_id: str,
         message: types.Content,
+        run_config: RunConfig | None = None,
         stream: bool = False,
     ):
         stream_mode = StreamingMode.SSE if stream else StreamingMode.NONE
 
-        async def event_generator():
-            async for event in self.runner.run_async(
-                user_id=self.user_id,
-                session_id=session_id,
-                new_message=message,
-                run_config=RunConfig(streaming_mode=stream_mode),
-            ):
-                if event.get_function_calls():
-                    for function_call in event.get_function_calls():
-                        logger.debug(f"Function call: {function_call}")
-                elif (
-                    event.content is not None
-                    and event.content.parts
-                    and event.content.parts[0].text is not None
-                    and len(event.content.parts[0].text.strip()) > 0
-                ):
-                    yield event.content.parts[0].text
+        if run_config is not None:
+            stream_mode = run_config.streaming_mode
+        else:
+            run_config = RunConfig(streaming_mode=stream_mode)
+        try:
 
-        final_output = ""
-        async for chunk in event_generator():
+            async def event_generator():
+                async for event in self.runner.run_async(
+                    user_id=self.user_id,
+                    session_id=session_id,
+                    new_message=message,
+                    run_config=run_config,
+                ):
+                    if event.get_function_calls():
+                        for function_call in event.get_function_calls():
+                            logger.debug(f"Function call: {function_call}")
+                    elif (
+                        event.content is not None
+                        and event.content.parts
+                        and event.content.parts[0].text is not None
+                        and len(event.content.parts[0].text.strip()) > 0
+                    ):
+                        yield event.content.parts[0].text
+
+            final_output = ""
+            async for chunk in event_generator():
+                if stream:
+                    print(chunk, end="", flush=True)
+                final_output += chunk
             if stream:
-                print(chunk, end="", flush=True)
-            final_output += chunk
-        if stream:
-            print()  # end with a new line
+                print()  # end with a new line
+        except LlmCallsLimitExceededError as e:
+            logger.warning(f"Max number of llm calls limit exceeded: {e}")
 
         return final_output
 
@@ -151,6 +166,7 @@ async def run(
         messages: RunnerMessage,
         session_id: str,
         stream: bool = False,
+        run_config: RunConfig | None = None,
         save_tracing_data: bool = False,
     ):
         converted_messages: list = self._convert_messages(messages)
@@ -163,7 +179,9 @@ async def run(
 
         final_output = ""
         for converted_message in converted_messages:
-            final_output = await self._run(session_id, converted_message, stream)
+            final_output = await self._run(
+                session_id, converted_message, run_config, stream
+            )
 
         # try to save tracing file
         if save_tracing_data:
@@ -193,6 +211,47 @@ def get_trace_id(self) -> str:
             logger.warning(f"Get tracer id failed as {e}")
             return "<unknown_trace_id>"
 
+    async def run_with_raw_message(
+        self,
+        message: types.Content,
+        session_id: str,
+        run_config: RunConfig | None = None,
+    ):
+        run_config = RunConfig() if not run_config else run_config
+
+        await self.short_term_memory.create_session(
+            app_name=self.app_name, user_id=self.user_id, session_id=session_id
+        )
+
+        try:
+
+            async def event_generator():
+                async for event in self.runner.run_async(
+                    user_id=self.user_id,
+                    session_id=session_id,
+                    new_message=message,
+                    run_config=run_config,
+                ):
+                    if event.get_function_calls():
+                        for function_call in event.get_function_calls():
+                            logger.debug(f"Function call: {function_call}")
+                    elif (
+                        event.content is not None
+                        and event.content.parts
+                        and event.content.parts[0].text is not None
+                        and len(event.content.parts[0].text.strip()) > 0
+                    ):
+                        yield event.content.parts[0].text
+
+            final_output = ""
+
+            async for chunk in event_generator():
+                final_output += chunk
+        except LlmCallsLimitExceededError as e:
+            logger.warning(f"Max number of llm calls limit exceeded: {e}")
+
+        return final_output
+
     def _print_trace_id(self) -> None:
         if not isinstance(self.agent, Agent):
             logger.warning(

veadk/tools/sandbox/browser_sandbox.py

@@ -12,16 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-browser_sandbox = ...
+from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
 
+from veadk.config import getenv
+from veadk.utils.mcp_utils import get_mcp_params
 
-def browser_use(prompt: str) -> str:
-    """Using the remote browser sandbox to according to the prompt.
+url = getenv("TOOL_BROWSER_SANDBOX_URL")
 
-    Args:
-        prompt (str): The prompt to be used.
 
-    Returns:
-        str: The response from the sandbox.
-    """
-    ...
+browser_sandbox = MCPToolset(connection_params=get_mcp_params(url=url))
+
+# browser_sandbox = ...
+
+
+# def browser_use(prompt: str) -> str:
+#     """Using the remote browser sandbox to according to the prompt.
+
+#     Args:
+#         prompt (str): The prompt to be used.
+
+#     Returns:
+#         str: The response from the sandbox.
+#     """
+#     ...

veadk/tools/sandbox/code_sandbox.py

@@ -12,19 +12,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-code_sandbox = ...
+from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
 
+from veadk.config import getenv
+from veadk.utils.mcp_utils import get_mcp_params
 
-def code_execution(code: str, language: str) -> str:
-    """Execute code in sandbox.
+url = getenv("TOOL_CODE_SANDBOX_URL")
 
-    Args:
-        code (str): The code to be executed.
-        language (str): The language of the code.
 
-    Returns:
-        str: The response from the sandbox.
-    """
+code_sandbox = MCPToolset(connection_params=get_mcp_params(url=url))
 
-    res = code_sandbox(code, language)
-    return res
+# code_sandbox = ...
+
+
+# def code_execution(code: str, language: str) -> str:
+#     """Execute code in sandbox.
+
+#     Args:
+#         code (str): The code to be executed.
+#         language (str): The language of the code.
+
+#     Returns:
+#         str: The response from the sandbox.
+#     """
+
+#     res = code_sandbox(code, language)
+#     return res

veadk/tracing/base_tracer.py

@@ -19,25 +19,6 @@
 logger = get_logger(__name__)
 
 
-def replace_bytes_with_empty(data):
-    """
-    Recursively traverse the data structure and replace all bytes types with empty strings.
-    Supports handling any nested structure of lists and dictionaries.
-    """
-    if isinstance(data, dict):
-        # Handle dictionary: Recursively process each value
-        return {k: replace_bytes_with_empty(v) for k, v in data.items()}
-    elif isinstance(data, list):
-        # Handle list: Recursively process each element
-        return [replace_bytes_with_empty(item) for item in data]
-    elif isinstance(data, bytes):
-        # When encountering the bytes type, replace it with an empty string
-        return "<image data>"
-    else:
-        # Keep other types unchanged
-        return data
-
-
 class BaseTracer(ABC):
     def __init__(self, name: str):
         self.name = name

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

@@ -371,13 +371,38 @@ def llm_output_value(params: LLMAttributesParams) -> ExtractorResponse:
     )
 
 
+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 =====
     "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,
+    "gen_ai.request.functions": llm_gen_ai_request_functions,
     # ===== response attributes =====
     "gen_ai.response.model": llm_gen_ai_response_model,
     "gen_ai.response.stop_reason": llm_gen_ai_response_stop_reason,
@@ -395,8 +420,8 @@ def llm_output_value(params: LLMAttributesParams) -> ExtractorResponse:
     # attributes
     "gen_ai.prompt": llm_gen_ai_prompt,
     "gen_ai.completion": llm_gen_ai_completion,
-    "input.value": llm_input_value,  # TLS required
-    "output.value": llm_output_value,  # TLS required
+    # "input.value": llm_input_value,
+    # "output.value": llm_output_value,
     # ===== usage =====
     "gen_ai.usage.input_tokens": llm_gen_ai_usage_input_tokens,
     "gen_ai.usage.output_tokens": llm_gen_ai_usage_output_tokens,

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

@@ -38,7 +38,7 @@ 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,
@@ -53,7 +53,7 @@ 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:
+def tool_gen_ai_tool_output(params: ToolAttributesParams) -> ExtractorResponse:
     function_response = params.function_response_event.get_function_responses()[
         0
     ].model_dump()
@@ -70,6 +70,8 @@ def tool_cozeloop_output(params: ToolAttributesParams) -> ExtractorResponse:
 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
 }