veadk-python 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

veadk/runner.py

@@ -11,32 +11,33 @@
 # 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.
+
 import asyncio
+import functools
+from types import MethodType
 from typing import Union
 
+from google import genai
 from google.adk.agents import RunConfig
+from google.adk.agents.base_agent import BaseAgent
 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
 from google.genai import types
 from google.genai.types import Blob
 
-from veadk.a2a.remote_ve_agent import RemoteVeAgent
 from veadk.agent import Agent
 from veadk.agents.loop_agent import LoopAgent
 from veadk.agents.parallel_agent import ParallelAgent
 from veadk.agents.sequential_agent import SequentialAgent
+from veadk.config import getenv
 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 read_png_to_bytes
-from veadk.integrations.ve_tos.ve_tos import VeTOS
+from veadk.utils.misc import formatted_timestamp, read_png_to_bytes
 
 logger = get_logger(__name__)
 
-
 RunnerMessage = Union[
     str,  # single turn text-based prompt
     list[str],  # multiple turn text-based prompt
@@ -45,111 +46,227 @@ RunnerMessage = Union[
     list[MediaMessage | str],  # multiple turn prompt with media and text-based prompt
 ]
 
-VeAgent = Union[Agent, RemoteVeAgent, SequentialAgent, ParallelAgent, LoopAgent]
 
+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
+            ):
+                process_func(
+                    part,
+                    self.app_name,
+                    user_id,
+                    session_id,
+                )
+    return
+
+
+def post_run_process(self):
+    return
+
+
+def intercept_new_message(process_func):
+    def decorator(func):
+        @functools.wraps(func)
+        async def wrapper(
+            self,
+            *,
+            user_id: str,
+            session_id: str,
+            new_message: types.Content,
+            **kwargs,
+        ):
+            pre_run_process(self, process_func, new_message, user_id, session_id)
+
+            async for event in func(
+                user_id=user_id,
+                session_id=session_id,
+                new_message=new_message,
+                **kwargs,
+            ):
+                yield event
+
+            post_run_process(self)
+
+        return wrapper
+
+    return decorator
+
+
+def _convert_messages(
+    messages: RunnerMessage,
+    app_name: str,
+    user_id: str,
+    session_id: str,
+) -> list:
+    """Convert VeADK formatted messages to Google ADK formatted 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."
+        )
+        _messages = [
+            types.Content(
+                role="user",
+                parts=[
+                    types.Part(text=messages.text),
+                    types.Part(
+                        inline_data=Blob(
+                            display_name=messages.media,
+                            data=read_png_to_bytes(messages.media),
+                            mime_type="image/png",
+                        )
+                    ),
+                ],
+            )
+        ]
+    elif isinstance(messages, list):
+        converted_messages = []
+        for message in messages:
+            converted_messages.extend(
+                _convert_messages(message, app_name, user_id, session_id)
+            )
+        _messages = converted_messages
+    else:
+        raise ValueError(f"Unknown message type: {type(messages)}")
+
+    return _messages
+
+
+def _upload_image_to_tos(
+    part: genai.types.Part, app_name: str, user_id: str, session_id: str
+) -> None:
+    try:
+        if part.inline_data and part.inline_data.display_name and part.inline_data.data:
+            from veadk.integrations.ve_tos.ve_tos import VeTOS
+
+            ve_tos = VeTOS()
+
+            object_key, tos_url = ve_tos.build_tos_url(
+                user_id=user_id,
+                app_name=app_name,
+                session_id=session_id,
+                data_path=part.inline_data.display_name,
+            )
+
+            upload_task = ve_tos.upload(object_key, part.inline_data.data)
+
+            if upload_task is not None:
+                asyncio.create_task(upload_task)
 
-class Runner:
+            part.inline_data.display_name = tos_url
+    except Exception as e:
+        logger.error(f"Upload to TOS failed: {e}")
+
+
+class Runner(ADKRunner):
     def __init__(
         self,
-        agent: VeAgent,
+        agent: BaseAgent | Agent,
         short_term_memory: ShortTermMemory | None = None,
-        plugins: list[BasePlugin] | None = None,
         app_name: str = "veadk_default_app",
         user_id: str = "veadk_default_user",
-    ):
-        self.app_name = app_name
+        upload_inline_data_to_tos: bool = False,
+        *args,
+        **kwargs,
+    ) -> None:
         self.user_id = user_id
+        self.long_term_memory = None
+        self.short_term_memory = short_term_memory
+        self.upload_inline_data_to_tos = upload_inline_data_to_tos
 
-        self.agent = agent
+        session_service = kwargs.pop("session_service", None)
+        memory_service = kwargs.pop("memory_service", None)
 
-        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
+        if session_service:
+            if short_term_memory:
+                logger.warning(
+                    "Short term memory is enabled, but session service is also provided. We will use session service from runner argument."
+                )
 
-        # prevent VeRemoteAgent has no long-term memory attr
-        if isinstance(self.agent, Agent):
-            self.long_term_memory = self.agent.long_term_memory
-        else:
-            self.long_term_memory = None
+        if not session_service:
+            if short_term_memory:
+                session_service = short_term_memory.session_service
+                logger.debug(
+                    f"Use session service {session_service} from short term memory."
+                )
+            else:
+                logger.warning(
+                    "No short term memory or session service provided, use an in-memory one instead."
+                )
+                short_term_memory = ShortTermMemory()
+                self.short_term_memory = short_term_memory
+                session_service = short_term_memory.session_service
+
+        if memory_service:
+            if hasattr(agent, "long_term_memory") and agent.long_term_memory:  # type: ignore
+                self.long_term_memory = agent.long_term_memory  # type: ignore
+                logger.warning(
+                    "Long term memory in agent is enabled, but memory service is also provided. We will use memory service from runner argument."
+                )
 
-        self.runner = ADKRunner(
-            app_name=self.app_name,
-            agent=self.agent,
-            session_service=self.session_service,
-            memory_service=self.long_term_memory,
-            plugins=plugins,
+        if not memory_service:
+            if hasattr(agent, "long_term_memory") and agent.long_term_memory:  # type: ignore
+                self.long_term_memory = agent.long_term_memory  # type: ignore
+                memory_service = agent.long_term_memory  # type: ignore
+            else:
+                logger.info("No long term memory provided.")
+
+        super().__init__(
+            agent=agent,
+            session_service=session_service,
+            memory_service=memory_service,
+            app_name=app_name,
+            *args,
+            **kwargs,
         )
 
-    def _convert_messages(self, messages, session_id) -> list:
-        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."
-            )
-            data = read_png_to_bytes(messages.media)
+        self.run_async = MethodType(
+            intercept_new_message(_upload_image_to_tos)(self.run_async), self
+        )
 
-            ve_tos = VeTOS()
-            object_key, tos_url = ve_tos.build_tos_url(
-                self.user_id, self.app_name, session_id, messages.media
-            )
-            try:
-                asyncio.create_task(ve_tos.upload(object_key, data))
-            except Exception as e:
-                logger.error(f"Upload to TOS failed: {e}")
-                tos_url = None
-
-            messages = [
-                types.Content(
-                    role="user",
-                    parts=[
-                        types.Part(text=messages.text),
-                        types.Part(
-                            inline_data=Blob(
-                                display_name=tos_url,
-                                data=data,
-                                mime_type="image/png",
-                            )
-                        ),
-                    ],
-                )
-            ]
-        elif isinstance(messages, list):
-            converted_messages = []
-            for message in messages:
-                converted_messages.extend(self._convert_messages(message, session_id))
-            messages = converted_messages
-        else:
-            raise ValueError(f"Unknown message type: {type(messages)}")
-
-        return messages
-
-    async def _run(
+    async def run(
         self,
-        session_id: str,
-        message: types.Content,
+        messages: RunnerMessage,
+        user_id: str = "",
+        session_id: str = f"tmp-session-{formatted_timestamp()}",
         run_config: RunConfig | None = None,
-        stream: bool = False,
+        save_tracing_data: bool = False,
+        upload_inline_data_to_tos: bool = False,
     ):
-        stream_mode = StreamingMode.SSE if stream else StreamingMode.NONE
+        if upload_inline_data_to_tos:
+            _upload_inline_data_to_tos = self.upload_inline_data_to_tos
+            self.upload_inline_data_to_tos = upload_inline_data_to_tos
+
+        if not run_config:
+            run_config = RunConfig(
+                # streaming_mode=stream_mode,
+                max_llm_calls=int(getenv("MODEL_AGENT_MAX_LLM_CALLS", 100)),
+            )
+        logger.info(f"Run config: {run_config}")
 
-        if run_config is not None:
-            stream_mode = run_config.streaming_mode
-        else:
-            run_config = RunConfig(streaming_mode=stream_mode)
-        try:
+        user_id = user_id or self.user_id
+
+        converted_messages: list = _convert_messages(
+            messages, self.app_name, user_id, session_id
+        )
+
+        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
+            )
 
-            async def event_generator():
-                async for event in self.runner.run_async(
+        final_output = ""
+        for converted_message in converted_messages:
+            try:
+                async for event in self.run_async(
                     user_id=self.user_id,
                     session_id=session_id,
-                    new_message=message,
+                    new_message=converted_message,
                     run_config=run_config,
                 ):
                     if event.get_function_calls():
@@ -161,41 +278,10 @@ class Runner:
                         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()  # end with a new line
-        except LlmCallsLimitExceededError as e:
-            logger.warning(f"Max number of llm calls limit exceeded: {e}")
-
-        return final_output
-
-    async def run(
-        self,
-        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, session_id)
-
-        await self.short_term_memory.create_session(
-            app_name=self.app_name, user_id=self.user_id, session_id=session_id
-        )
-
-        logger.info("Begin to process user messages.")
-
-        final_output = ""
-        for converted_message in converted_messages:
-            final_output = await self._run(
-                session_id, converted_message, run_config, stream
-            )
+                        final_output += event.content.parts[0].text
+            except LlmCallsLimitExceededError as e:
+                logger.warning(f"Max number of llm calls limit exceeded: {e}")
+                final_output = ""
 
         # try to save tracing file
         if save_tracing_data:
@@ -203,6 +289,9 @@ class Runner:
 
         self._print_trace_id()
 
+        if upload_inline_data_to_tos:
+            self.upload_inline_data_to_tos = _upload_inline_data_to_tos  # type: ignore
+
         return final_output
 
     def get_trace_id(self) -> str:
@@ -225,47 +314,6 @@ class Runner:
             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(
@@ -336,61 +384,3 @@ class Runner:
 
         await self.long_term_memory.add_session_to_memory(session)
         logger.info(f"Add session `{session.id}` to long term memory.")
-
-    # [deprecated] we will not host a chat-service in VeADK, so the following two methods are deprecated
-
-    # async def run_with_final_event(
-    #     self,
-    #     messages: RunnerMessage,
-    #     session_id: str,
-    # ):
-    #     """non-streaming run with final event"""
-    #     messages: list = self._convert_messages(messages)
-
-    #     await self.short_term_memory.create_session(
-    #         app_name=self.app_name, user_id=self.user_id, session_id=session_id
-    #     )
-
-    #     logger.info("Begin to process user messages.")
-
-    #     final_event = ""
-    #     async for event in self.runner.run_async(
-    #         user_id=self.user_id, session_id=session_id, new_message=messages[0]
-    #     ):
-    #         if event.get_function_calls():
-    #             for function_call in event.get_function_calls():
-    #                 logger.debug(f"Function call: {function_call}")
-    #         elif (
-    #             not event.partial
-    #             and event.content.parts[0].text is not None
-    #             and len(event.content.parts[0].text.strip()) > 0
-    #         ):
-    #             final_event = event.model_dump_json(exclude_none=True, by_alias=True)
-
-    #     return final_event
-
-    # async def run_sse(
-    #     self,
-    #     session_id: str,
-    #     prompt: str,
-    # ):
-    #     message = types.Content(role="user", parts=[types.Part(text=prompt)])
-
-    #     await self.short_term_memory.create_session(
-    #         app_name=self.app_name, user_id=self.user_id, session_id=session_id
-    #     )
-
-    #     logger.info("Begin to process user messages under SSE method.")
-
-    #     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=StreamingMode.SSE),
-    #     ):
-    #         # Format as SSE data
-    #         sse_event = event.model_dump_json(exclude_none=True, by_alias=True)
-    #         if event.get_function_calls():
-    #             for function_call in event.get_function_calls():
-    #                 logger.debug(f"SSE function call event: {sse_event}")
-    #         yield f"data: {sse_event}\n\n"

veadk/tools/builtin_tools/vesearch.py

@@ -14,7 +14,7 @@
 
 import requests
 
-from veadk.config import getenv
+from veadk.config import getenv, settings
 
 
 def vesearch(query: str) -> str:
@@ -26,7 +26,7 @@ def vesearch(query: str) -> str:
     Returns:
         Summarized search results.
     """
-    api_key = getenv("TOOL_VESEARCH_API_KEY")
+    api_key = settings.tool.vesearch.api_key
     bot_id = str(getenv("TOOL_VESEARCH_ENDPOINT"))
 
     if api_key == "":

veadk/tools/builtin_tools/video_generate.py

@@ -12,18 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict
-from google.adk.tools import ToolContext
-from volcenginesdkarkruntime import Ark
-from veadk.config import getenv
+import json
 import time
 import traceback
-import json
-from veadk.version import VERSION
+from typing import Dict, cast
+
+from google.adk.tools import ToolContext
 from opentelemetry import trace
 from opentelemetry.trace import Span
+from volcenginesdkarkruntime import Ark
+from volcenginesdkarkruntime.types.content_generation.create_task_content_param import (
+    CreateTaskContentParam,
+)
 
+from veadk.config import getenv
 from veadk.utils.logger import get_logger
+from veadk.version import VERSION
 
 logger = get_logger(__name__)
 
@@ -47,13 +51,16 @@ async def generate(prompt, first_frame_image=None, last_frame_image=None):
             logger.debug("first frame generation")
             response = client.content_generation.tasks.create(
                 model=getenv("MODEL_VIDEO_NAME"),
-                content=[
-                    {"type": "text", "text": prompt},
-                    {
-                        "type": "image_url",
-                        "image_url": {"url": first_frame_image},
-                    },
-                ],
+                content=cast(
+                    list[CreateTaskContentParam],  # avoid IDE warning
+                    [
+                        {"type": "text", "text": prompt},
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": first_frame_image},
+                        },
+                    ],
+                ),
             )
         else:
             logger.debug("last frame generation")
@@ -277,13 +284,13 @@ async def video_generate(params: list, tool_context: ToolContext) -> Dict:
 def add_span_attributes(
     span: Span,
     tool_context: ToolContext,
-    input_part: dict = None,
-    output_part: dict = None,
-    input_tokens: int = None,
-    output_tokens: int = None,
-    total_tokens: int = None,
-    request_model: str = None,
-    response_model: str = None,
+    input_part: dict | None = None,
+    output_part: dict | None = None,
+    input_tokens: int | None = None,
+    output_tokens: int | None = None,
+    total_tokens: int | None = None,
+    request_model: str | None = None,
+    response_model: str | None = None,
 ):
     try:
         # common attributes

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

@@ -49,6 +49,10 @@ def common_cozeloop_report_source(**kwargs) -> str:
     return "veadk"
 
 
+def common_cozeloop_call_type(**kwargs) -> str:
+    return kwargs.get("call_type")
+
+
 def llm_openinference_instrumentation_veadk(**kwargs) -> str:
     return VERSION
 
@@ -68,4 +72,5 @@ COMMON_ATTRIBUTES = {
     "user.id": common_gen_ai_user_id,  # CozeLoop / TLS required
     "session.id": common_gen_ai_session_id,  # CozeLoop / TLS required
     "cozeloop.report.source": common_cozeloop_report_source,  # CozeLoop required
+    "cozeloop.call_type": common_cozeloop_call_type,  # CozeLoop required
 }