veadk-python 0.1.0__py3-none-any.whl

veadk/__init__.py

@@ -0,0 +1,31 @@
+# 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 .version import VERSION
+
+
+# Lazy loading for `Agent` class
+def __getattr__(name):
+    if name == "Agent":
+        from .agent import Agent
+
+        return Agent
+    if name == "Runner":
+        from .runner import Runner
+
+        return Runner
+    raise AttributeError(f"module 'veadk' has no attribute '{name}'")
+
+
+__all__ = ["Agent", "Runner", "VERSION"]

veadk/a2a/__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/a2a/agent_card.py

@@ -0,0 +1,45 @@
+# 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 a2a.types import AgentCapabilities, AgentCard, AgentProvider, AgentSkill
+
+from veadk import Agent
+from veadk.version import VERSION
+
+
+def get_agent_card(
+    agent: Agent, url: str, version: str = VERSION, provider: str = "veadk"
+) -> AgentCard:
+    agent_provider = AgentProvider(organization=provider, url="")
+    agent_capabilities = AgentCapabilities()
+    agent_skills = [
+        AgentSkill(
+            id="0",
+            name="chat",
+            description="Basically chat with user.",
+            tags=["chat", "talk"],
+        )
+    ]
+    agent_card = AgentCard(
+        capabilities=agent_capabilities,
+        description=agent.description,
+        name=agent.name,
+        defaultInputModes=["text"],
+        defaultOutputModes=["text"],
+        provider=agent_provider,
+        skills=agent_skills,
+        url=url,
+        version=version,
+    )
+    return agent_card

veadk/a2a/remote_ve_agent.py

@@ -0,0 +1,19 @@
+import json
+
+import requests
+from a2a.types import AgentCard
+from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
+
+AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent.json"
+
+
+class RemoteVeAgent(RemoteA2aAgent):
+    def __init__(self, name: str, url: str):
+        agent_card_dict = requests.get(url + AGENT_CARD_WELL_KNOWN_PATH).json()
+        agent_card_dict["url"] = url
+
+        agent_card_json_str = json.dumps(agent_card_dict, ensure_ascii=False, indent=2)
+
+        agent_card_object = AgentCard.model_validate_json(str(agent_card_json_str))
+
+        super().__init__(name=name, agent_card=agent_card_object)

veadk/a2a/ve_a2a_server.py

@@ -0,0 +1,77 @@
+# 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 a2a.server.apps.jsonrpc.fastapi_app import A2AFastAPIApplication
+from a2a.server.request_handlers.default_request_handler import DefaultRequestHandler
+from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
+from fastapi import FastAPI
+
+from veadk import Agent
+from veadk.a2a.agent_card import get_agent_card
+from veadk.a2a.ve_agent_executor import VeAgentExecutor
+from veadk.memory.short_term_memory import ShortTermMemory
+
+
+class VeA2AServer:
+    def __init__(
+        self, agent: Agent, url: str, app_name: str, short_term_memory: ShortTermMemory
+    ):
+        self.agent_card = get_agent_card(agent, url)
+
+        self.agent_executor = VeAgentExecutor(
+            app_name=app_name,
+            agent=agent,
+            short_term_memory=short_term_memory,
+        )
+
+        self.task_store = InMemoryTaskStore()
+
+        self.request_handler = DefaultRequestHandler(
+            agent_executor=self.agent_executor, task_store=self.task_store
+        )
+
+    def build(self) -> FastAPI:
+        app_application = A2AFastAPIApplication(
+            agent_card=self.agent_card,
+            http_handler=self.request_handler,
+        )
+        app = app_application.build()  # build routes
+
+        # import uvicorn
+        # uvicorn.run(app, host="127.0.0.1", port=8000)
+        return app
+
+
+def init_app(
+    server_url: str, app_name: str, agent: Agent, short_term_memory: ShortTermMemory
+) -> FastAPI:
+    """Init the fastapi application in terms of VeADK agent.
+
+    Args:
+        server_url: str, the url of the server
+        app_name: str, the name of the app
+        agent: Agent, the agent of the app
+        short_term_memory: ShortTermMemory, the short term memory of the app
+
+    Returns:
+        FastAPI, the fastapi app
+    """
+
+    server = VeA2AServer(
+        agent=agent,
+        url=server_url,
+        app_name=app_name,
+        short_term_memory=short_term_memory,
+    )
+    return server.build()

veadk/a2a/ve_agent_executor.py

@@ -0,0 +1,78 @@
+# 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 a2a.server.agent_execution.agent_executor import AgentExecutor
+from a2a.server.agent_execution.context import RequestContext
+from a2a.server.events.event_queue import EventQueue
+from a2a.utils import new_agent_text_message
+from typing_extensions import override
+
+from veadk import Agent
+from veadk.memory.short_term_memory import ShortTermMemory
+from veadk.runner import Runner
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class VeAgentExecutor(AgentExecutor):
+    def __init__(self, app_name: str, agent: Agent, short_term_memory: ShortTermMemory):
+        super().__init__()
+        self.app_name = app_name
+        self.agent = agent
+        self.short_term_memory = short_term_memory
+
+        self.runner = Runner(
+            agent=self.agent,
+            short_term_memory=self.short_term_memory,
+            app_name=app_name,
+            user_id="",  # waiting for `execute` change
+        )
+
+    @override
+    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
+        # extract metadata
+        user_id = (
+            context.metadata["user_id"]
+            if "user_id" in context.metadata
+            else "unkonwn_user"
+        )
+        self.runner.user_id = user_id
+
+        session_id = (
+            context.metadata["session_id"]
+            if "session_id" in context.metadata
+            else "unkonwn_session"
+        )
+
+        # process user input
+        user_input = context.get_user_input()
+
+        logger.debug(
+            f"Request: user_id: {user_id}, session_id: {session_id}, user_input: {user_input}"
+        )
+
+        # running
+        final_output = await self.runner.run(
+            messages=user_input,
+            session_id=session_id,
+        )
+
+        logger.debug(f"Final output: {final_output}")
+
+        await event_queue.enqueue_event(new_agent_text_message(final_output))
+
+    @override
+    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
+        return None

veadk/a2a/ve_task_store.py

@@ -0,0 +1,37 @@
+# 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 a2a.server.tasks import TaskStore
+from a2a.types import Task
+from typing_extensions import override
+
+
+class VeTaskStore(TaskStore):
+    def __init__(self):
+        super().__init__()
+
+    @override
+    async def save(self, task: Task) -> None:
+        """Saves or updates a task in the store."""
+        return None
+
+    @override
+    async def get(self, task_id: str) -> Task | None:
+        """Retrieves a task from the store by ID."""
+        return None
+
+    @override
+    async def delete(self, task_id: str) -> None:
+        """Deletes a task from the store by ID."""
+        return None

veadk/agent.py

@@ -0,0 +1,253 @@
+# 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 typing import Optional
+
+from google.adk.agents import LlmAgent, RunConfig
+from google.adk.agents.llm_agent import ToolUnion
+from google.adk.agents.run_config import StreamingMode
+from google.adk.models.lite_llm import LiteLlm
+from google.adk.runners import Runner
+from google.genai import types
+from pydantic import ConfigDict, Field
+from typing_extensions import Any
+
+from veadk.config import getenv
+from veadk.consts import (
+    DEFALUT_MODEL_AGENT_PROVIDER,
+    DEFAULT_MODEL_AGENT_API_BASE,
+    DEFAULT_MODEL_AGENT_NAME,
+)
+from veadk.evaluation import EvalSetRecorder
+from veadk.knowledgebase import KnowledgeBase
+from veadk.memory.long_term_memory import LongTermMemory
+from veadk.memory.short_term_memory import ShortTermMemory
+from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION
+from veadk.tracing.base_tracer import BaseTracer
+from veadk.utils.logger import get_logger
+from veadk.utils.patches import patch_asyncio
+
+patch_asyncio()
+logger = get_logger(__name__)
+
+
+class Agent(LlmAgent):
+    """LLM-based Agent with Volcengine capabilities."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
+    """The model config"""
+
+    name: str = "veAgent"
+    """The name of the agent."""
+
+    description: str = DEFAULT_DESCRIPTION
+    """The description of the agent. This will be helpful in A2A scenario."""
+
+    instruction: str = DEFAULT_INSTRUCTION
+    """The instruction for the agent, such as principles of function calling."""
+
+    # factory
+    model_name: str = getenv("MODEL_AGENT_NAME", DEFAULT_MODEL_AGENT_NAME)
+    """The name of the model for agent running."""
+
+    model_provider: str = getenv("MODEL_AGENT_PROVIDER", DEFALUT_MODEL_AGENT_PROVIDER)
+    """The provider of the model for agent running."""
+
+    model_api_base: str = getenv("MODEL_AGENT_API_BASE", DEFAULT_MODEL_AGENT_API_BASE)
+    """The api base of the model for agent running."""
+
+    model_api_key: str = Field(
+        ..., default_factory=lambda: getenv("MODEL_AGENT_API_KEY")
+    )
+    """The api key of the model for agent running."""
+
+    tools: list[ToolUnion] = []
+    """The tools provided to agent."""
+
+    sub_agents: list[Agent] = Field(default_factory=list, exclude=True)
+    """The sub agents provided to agent."""
+
+    knowledgebase: Optional[KnowledgeBase] = None
+    """The knowledgebase provided to agent."""
+
+    long_term_memory: Optional[LongTermMemory] = None
+    """The long term memory provided to agent.
+
+    In VeADK, the `long_term_memory` refers to cross-session memory under the same user.
+    """
+
+    tracers: list[BaseTracer] = []
+    """The tracers provided to agent."""
+
+    serve_url: str = ""
+    """The url of agent serving host. Show in agent card."""
+
+    def model_post_init(self, __context: Any) -> None:
+        super().model_post_init(None)  # for sub_agents init
+        self.model = LiteLlm(
+            model=f"{self.model_provider}/{self.model_name}",
+            api_key=self.model_api_key,
+            api_base=self.model_api_base,
+        )
+
+        if self.knowledgebase:
+            from veadk.tools import load_knowledgebase_tool
+
+            load_knowledgebase_tool.knowledgebase = self.knowledgebase
+            self.tools.append(load_knowledgebase_tool.load_knowledgebase_tool)
+
+        if self.long_term_memory is not None:
+            from google.adk.tools import load_memory
+
+            self.tools.append(load_memory)
+
+        if self.tracers:
+            self.before_model_callback = []
+            self.after_model_callback = []
+            for tracer in self.tracers:
+                self.before_model_callback.append(tracer.llm_metrics_hook)
+                self.after_model_callback.append(tracer.token_metrics_hook)
+
+        logger.info(f"Agent `{self.name}` init done.")
+        logger.debug(
+            f"Agent: {self.model_dump(include={'name', 'model_name', 'model_api_base', 'tools', 'serve_url'})}"
+        )
+
+    async def _run(
+        self,
+        runner,
+        user_id: str,
+        session_id: str,
+        message: types.Content,
+        stream: bool,
+    ):
+        stream_mode = StreamingMode.SSE if stream else StreamingMode.NONE
+
+        async def event_generator():
+            async for event in runner.run_async(
+                user_id=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,
+        prompt: str | list[str],
+        stream: bool = False,
+        app_name: str = "veadk_app",
+        user_id: str = "veadk_user",
+        session_id="veadk_session",
+        load_history_sessions_from_db: bool = False,
+        db_url: str = "",
+        collect_runtime_data: bool = False,
+        eval_set_id: str = "",
+        save_session_to_memory: bool = False,
+        enable_memory_optimization: bool = False,
+    ):
+        """Running the agent. The runner and session service will be created automatically.
+
+        For production, consider using Google-ADK runner to run agent, rather than invoking this method.
+
+        Args:
+            prompt (str | list[str]): The prompt to run the agent.
+            stream (bool, optional): Whether to stream the output. Defaults to False.
+            app_name (str, optional): The name of the application. Defaults to "veadk_app".
+            user_id (str, optional): The id of the user. Defaults to "veadk_user".
+            session_id (str, optional): The id of the session. Defaults to "veadk_session".
+            load_history_sessions_from_db (bool, optional): Whether to load history sessions from database. Defaults to False.
+            db_url (str, optional): The url of the database. Defaults to "".
+            collect_runtime_data (bool, optional): Whether to collect runtime data. Defaults to False.
+            eval_set_id (str, optional): The id of the eval set. Defaults to "".
+            save_session_to_memory (bool, optional): Whether to save this turn session to memory. Defaults to False.
+        """
+
+        logger.warning(
+            "Running agent in this function is only for development and testing, do not use this function in production. For production, consider using `Google ADK Runner` to run agent, rather than invoking this method."
+        )
+        logger.info(
+            f"Run agent {self.name}: app_name: {app_name}, user_id: {user_id}, session_id: {session_id}."
+        )
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        # memory service
+        short_term_memory = ShortTermMemory(
+            backend="database" if load_history_sessions_from_db else "local",
+            enable_memory_optimization=enable_memory_optimization,
+            db_url=db_url,
+        )
+        session_service = short_term_memory.session_service
+        await short_term_memory.create_session(
+            app_name=app_name, user_id=user_id, session_id=session_id
+        )
+
+        # runner
+        runner = Runner(
+            agent=self,
+            app_name=app_name,
+            session_service=session_service,
+            memory_service=self.long_term_memory,
+        )
+
+        logger.info(f"Begin to process prompt {prompt}")
+        # run
+        final_output = ""
+        for _prompt in prompt:
+            message = types.Content(role="user", parts=[types.Part(text=_prompt)])
+            final_output = await self._run(runner, user_id, session_id, message, stream)
+
+        # VeADK features
+        if save_session_to_memory:
+            assert self.long_term_memory is not None, (
+                "Long-term memory is not initialized in agent"
+            )
+            session = await session_service.get_session(
+                app_name=app_name,
+                user_id=user_id,
+                session_id=session_id,
+            )
+            await self.long_term_memory.add_session_to_memory(session)
+            logger.info(f"Add session `{session.id}` to your long-term memory.")
+
+        if collect_runtime_data:
+            eval_set_recorder = EvalSetRecorder(session_service, eval_set_id)
+            dump_path = await eval_set_recorder.dump(app_name, user_id, session_id)
+            self._dump_path = dump_path  # just for test/debug/instrumentation
+
+        if self.tracers:
+            for tracer in self.tracers:
+                tracer.dump(user_id, session_id)
+
+        return final_output

veadk/cli/__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.