veadk-python 0.1.0__py3-none-any.whl

veadk/runner.py

@@ -0,0 +1,252 @@
+# 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 typing import Union
+
+from google.adk.agents import RunConfig
+from google.adk.agents.run_config import StreamingMode
+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.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
+
+logger = get_logger(__name__)
+
+
+RunnerMessage = Union[
+    str,  # single turn text-based prompt
+    list[str],  # multiple turn text-based prompt
+    MediaMessage,  # single turn prompt with media
+    list[MediaMessage],  # multiple turn prompt with media
+    list[MediaMessage | str],  # multiple turn prompt with media and text-based prompt
+]
+
+
+class Runner:
+    def __init__(
+        self,
+        agent: Agent | RemoteVeAgent,
+        short_term_memory: ShortTermMemory,
+        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 isinstance(self.agent, Agent):
+            self.long_term_memory = self.agent.long_term_memory
+        else:
+            self.long_term_memory = None
+
+        # maintain a in-memory runner for fast inference
+        self.runner = ADKRunner(
+            app_name=self.app_name,
+            agent=self.agent,
+            session_service=self.session_service,
+            memory_service=self.long_term_memory,
+        )
+
+    def _convert_messages(self, messages) -> 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."
+            )
+            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(self._convert_messages(message))
+            messages = converted_messages
+        else:
+            raise ValueError(f"Unknown message type: {type(messages)}")
+
+        return messages
+
+    async def _run(
+        self,
+        session_id: str,
+        message: types.Content,
+        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[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
+
+        return final_output
+
+    async def run(
+        self,
+        messages: RunnerMessage,
+        session_id: str,
+        stream: bool = False,
+    ):
+        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_output = ""
+        for message in messages:
+            final_output = await self._run(session_id, message, stream)
+
+        # try to save tracing file
+        if isinstance(self.agent, Agent):
+            self.save_tracing_file(session_id)
+
+        return final_output
+
+    def save_tracing_file(self, session_id: str) -> str:
+        if not self.agent.tracers:
+            return
+
+        try:
+            dump_path = ""
+            for tracer in self.agent.tracers:
+                dump_path = tracer.dump(self.user_id, session_id)
+
+            return dump_path
+        except Exception as e:
+            logger.error(f"Failed to save tracing file: {e}")
+            return ""
+
+    async def save_eval_set(self, session_id: str, eval_set_id: str = "default") -> str:
+        eval_set_recorder = EvalSetRecorder(self.session_service, eval_set_id)
+        eval_set_path = await eval_set_recorder.dump(
+            self.app_name, self.user_id, session_id
+        )
+        return eval_set_path
+
+    async def save_session_to_long_term_memory(self, session_id: str) -> None:
+        if not self.long_term_memory:
+            logger.warning("Long-term memory is not enabled. Failed to save session.")
+            return
+
+        session = await self.session_service.get_session(
+            app_name=self.app_name,
+            user_id=self.user_id,
+            session_id=session_id,
+        )
+        await self.long_term_memory.add_session_to_memory(session)
+        logger.info(f"Add session `{session.id}` to long-term memory.")
+
+    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/__init__.py

@@ -0,0 +1,13 @@
+# 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.

veadk/tools/builtin_tools/__init__.py

@@ -0,0 +1,13 @@
+# 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.

veadk/tools/builtin_tools/lark.py

@@ -0,0 +1,67 @@
+# 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.
+
+import subprocess
+
+from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset, StdioServerParameters
+
+from veadk.config import getenv
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def check_env():
+    try:
+        result = subprocess.run(
+            ["npx", "-v"], capture_output=True, text=True, check=True
+        )
+        version = result.stdout.strip()
+        logger.info(f"Check `npx` command done, version: {version}")
+    except Exception as e:
+        raise Exception(
+            "Check `npx` command failed. Please install `npx` command manually."
+        ) from e
+
+    try:
+        result = subprocess.run(
+            ["lark-mcp", "-V"], capture_output=True, text=True, check=True
+        )
+        version = result.stdout.strip()
+        logger.info(f"Check `lark-mcp` command done, version: {version}")
+    except Exception as e:
+        raise Exception(
+            "Check `lark-mcp` command failed. Please install it manually."
+        ) from e
+
+
+check_env()
+
+lark_tools = MCPToolset(
+    connection_params=StdioServerParameters(
+        command="npx",
+        args=[
+            "-y",
+            "@larksuiteoapi/lark-mcp",
+            "mcp",
+            "-a",
+            getenv("TOOL_LARK_ENDPOINT"),
+            "-s",
+            getenv("TOOL_LARK_API_KEY"),
+            "-u",
+            getenv("TOOL_LARK_TOKEN"),
+        ],
+        errlog=None,
+    ),
+)

veadk/tools/builtin_tools/las.py

@@ -0,0 +1,23 @@
+# 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 google.adk.tools.mcp_tool.mcp_toolset import MCPToolset, SseConnectionParams
+
+from veadk.config import getenv
+
+url = getenv("TOOL_LAS_URL")
+dataset_id = getenv("TOOL_LAS_DATASET_ID")
+
+
+las = MCPToolset(connection_params=SseConnectionParams(url=url))

veadk/tools/builtin_tools/vesearch.py

@@ -0,0 +1,49 @@
+# 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.
+
+import requests
+
+from veadk.config import getenv
+
+
+def vesearch(query: str) -> str:
+    """Search information from Internet, social media, news sites, etc.
+
+    Args:
+        query: The query string to search.
+
+    Returns:
+        Summarized search results.
+    """
+    api_key = getenv("TOOL_VESEARCH_API_KEY")
+    bot_id = str(getenv("TOOL_VESEARCH_ENDPOINT"))
+
+    if api_key == "":
+        return "Invoke `vesearch` failed. Please set VESEARCH_API_KEY as your environment variable."
+    if bot_id == "":
+        return "Invoke `vesearch` failed. Please set VESEARCH_BOT_ID as your environment variable."
+
+    URL = "https://open.feedcoopapi.com/agent_api/agent/chat/completion"
+    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+    data = {
+        "bot_id": bot_id,
+        "messages": [{"role": "user", "content": query}],
+    }
+
+    response = requests.post(URL, json=data, headers=headers)
+    if response.status_code == 200:
+        result = response.json()
+        return result["choices"][0]["message"]["content"]
+    else:
+        return response.text

veadk/tools/builtin_tools/web_scraper.py

@@ -0,0 +1,76 @@
+# 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 typing import Any
+
+import requests
+
+from veadk.config import getenv
+
+
+def web_scraper(query: str) -> dict[str, Any]:
+    """query a single keyword from some search engineer
+
+    Args:
+        query (str, optional): query keyword. Defaults to None.
+
+    Returns:
+        dict[str, Any]: search results
+    """
+
+    endpoint = getenv("TOOL_WEB_SCRAPER_ENDPOINT")
+    token = getenv("TOOL_WEB_SCRAPER_API_KEY")
+
+    try:
+        url = f"https://{endpoint}/v1/queries"
+        headers = {
+            "Content-Type": "application/json",
+            "X-VE-Source": "google_search",
+            "X-VE-API-Key": token,
+        }
+        data = {
+            "query": query,
+            "source": "google_search",
+            "parse": True,
+            "limit": "10",
+            "start_page": "1",
+            "pages": "1",
+            "context": [
+                {"key": "nfpr", "value": True},
+                {"key": "safe_search", "value": False},
+                {"key": "filter", "value": 1},
+            ],
+        }
+        response = requests.post(url, headers=headers, json=data, verify=False)
+
+        response.raise_for_status()
+
+        response_json = response.json() if response.content else None
+        results_dict = (
+            response_json.get("results")[0].get("content").get("results").get("organic")
+        )
+
+        results_str = ""
+        for r in results_dict:
+            dic = {
+                "url": r.get("url"),
+                "title": r.get("title"),
+                "description": r.get("desc"),
+            }
+            results_str += str(dic)
+        return results_str
+
+    except requests.exceptions.RequestException as e:
+        error_message = f"Error: {str(e)}"
+        raise ValueError(error_message)

veadk/tools/builtin_tools/web_search.py

@@ -0,0 +1,192 @@
+# 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.
+
+"""
+The document of this tool see: https://www.volcengine.com/docs/85508/1650263
+"""
+
+import datetime
+import hashlib
+import hmac
+import json
+from urllib.parse import quote
+
+import requests
+
+from veadk.config import getenv
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+Service = "volc_torchlight_api"
+Version = "2025-01-01"
+Region = "cn-beijing"
+Host = "mercury.volcengineapi.com"
+ContentType = "application/json"
+
+
+def norm_query(params):
+    query = ""
+    for key in sorted(params.keys()):
+        if isinstance(params[key], list):
+            for k in params[key]:
+                query = (
+                    query + quote(key, safe="-_.~") + "=" + quote(k, safe="-_.~") + "&"
+                )
+        else:
+            query = (
+                query
+                + quote(key, safe="-_.~")
+                + "="
+                + quote(params[key], safe="-_.~")
+                + "&"
+            )
+    query = query[:-1]
+    return query.replace("+", "%20")
+
+
+def hmac_sha256(key: bytes, content: str):
+    return hmac.new(key, content.encode("utf-8"), hashlib.sha256).digest()
+
+
+def hash_sha256(content: str):
+    return hashlib.sha256(content.encode("utf-8")).hexdigest()
+
+
+def request(method, date, query, header, ak, sk, action, body):
+    credential = {
+        "access_key_id": ak,
+        "secret_access_key": sk,
+        "service": Service,
+        "region": Region,
+    }
+    request_param = {
+        "body": body,
+        "host": Host,
+        "path": "/",
+        "method": method,
+        "content_type": ContentType,
+        "date": date,
+        "query": {"Action": action, "Version": Version, **query},
+    }
+    if body is None:
+        request_param["body"] = ""
+    # 第四步：接下来开始计算签名。在计算签名前，先准备好用于接收签算结果的 signResult 变量，并设置一些参数。
+    # 初始化签名结果的结构体
+    x_date = request_param["date"].strftime("%Y%m%dT%H%M%SZ")
+    short_x_date = x_date[:8]
+    x_content_sha256 = hash_sha256(request_param["body"])
+    sign_result = {
+        "Host": request_param["host"],
+        "X-Content-Sha256": x_content_sha256,
+        "X-Date": x_date,
+        "Content-Type": request_param["content_type"],
+    }
+    signed_headers_str = ";".join(
+        ["content-type", "host", "x-content-sha256", "x-date"]
+    )
+    # signed_headers_str = signed_headers_str + ";x-security-token"
+    canonical_request_str = "\n".join(
+        [
+            request_param["method"].upper(),
+            request_param["path"],
+            norm_query(request_param["query"]),
+            "\n".join(
+                [
+                    "content-type:" + request_param["content_type"],
+                    "host:" + request_param["host"],
+                    "x-content-sha256:" + x_content_sha256,
+                    "x-date:" + x_date,
+                ]
+            ),
+            "",
+            signed_headers_str,
+            x_content_sha256,
+        ]
+    )
+
+    hashed_canonical_request = hash_sha256(canonical_request_str)
+
+    credential_scope = "/".join(
+        [short_x_date, credential["region"], credential["service"], "request"]
+    )
+    string_to_sign = "\n".join(
+        ["HMAC-SHA256", x_date, credential_scope, hashed_canonical_request]
+    )
+
+    k_date = hmac_sha256(credential["secret_access_key"].encode("utf-8"), short_x_date)
+    k_region = hmac_sha256(k_date, credential["region"])
+    k_service = hmac_sha256(k_region, credential["service"])
+    k_signing = hmac_sha256(k_service, "request")
+    signature = hmac_sha256(k_signing, string_to_sign).hex()
+
+    sign_result["Authorization"] = (
+        "HMAC-SHA256 Credential={}, SignedHeaders={}, Signature={}".format(
+            credential["access_key_id"] + "/" + credential_scope,
+            signed_headers_str,
+            signature,
+        )
+    )
+    header = {**header, **sign_result}
+    # header = {**header, **{"X-Security-Token": SessionToken}}
+    r = requests.request(
+        method=method,
+        url="https://{}{}".format(request_param["host"], request_param["path"]),
+        headers=header,
+        params=request_param["query"],
+        data=request_param["body"],
+    )
+    return r.json()
+
+
+def web_search(query: str) -> list[str]:
+    """Search a query in websites.
+
+    Args:
+        query: The query to search.
+
+    Returns:
+        A list of result documents.
+    """
+    req = {
+        "Query": query,
+        "SearchType": "web",
+        "Count": 5,
+        "NeedSummary": True,
+    }
+    ak = getenv("VOLCENGINE_ACCESS_KEY")
+    sk = getenv("VOLCENGINE_SECRET_KEY")
+
+    now = datetime.datetime.utcnow()
+    response_body = request(
+        "POST",
+        now,
+        {},
+        {},
+        ak,
+        sk,
+        "WebSearch",
+        json.dumps(req),
+    )
+    try:
+        results: list = response_body["Result"]["WebResults"]
+    except Exception as e:
+        logger.error(f"Web search failed: {e}")
+        return []
+
+    final_results = []
+    for result in results:
+        final_results.append(result["Summary"].strip())
+    return final_results