veadk-python 0.2.27__py3-none-any.whl

veadk/memory/short_term_memory_processor.py

@@ -0,0 +1,100 @@
+# 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 functools
+import json
+
+from google.adk.events.event import Event
+from google.adk.sessions import Session
+from google.genai.types import Content, Part
+from litellm import completion
+
+from veadk.config import settings
+from veadk.consts import (
+    DEFAULT_MODEL_AGENT_API_BASE,
+    DEFAULT_MODEL_AGENT_NAME,
+    DEFAULT_MODEL_AGENT_PROVIDER,
+)
+from veadk.prompts.prompt_memory_processor import render_prompt
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class ShortTermMemoryProcessor:
+    def __init__(self) -> None: ...
+
+    def patch(self):
+        """Patch the `get_session` function"""
+
+        def intercept_get_session(func):
+            @functools.wraps(func)
+            async def wrapper(*args, **kwargs):
+                session = await func(*args, **kwargs)
+                if session:
+                    abstracted_session = self.after_load_session(session)
+                else:
+                    abstracted_session = session
+                return abstracted_session
+
+            return wrapper
+
+        return intercept_get_session
+
+    def after_load_session(self, session: Session) -> Session:
+        messages = []
+        for event in session.events:
+            content = event.content
+            if not content or not content.parts:
+                continue
+            message = {
+                "role": content.role,
+                "content": content.parts[0].text,
+            }
+            messages.append(message)
+
+        prompt = render_prompt(messages=messages)
+
+        res = completion(
+            model=DEFAULT_MODEL_AGENT_PROVIDER + "/" + DEFAULT_MODEL_AGENT_NAME,
+            base_url=DEFAULT_MODEL_AGENT_API_BASE,
+            api_key=settings.model.api_key,
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt,
+                }
+            ],
+        )
+        logger.debug(f"Response from memory optimization model: {res}")
+
+        extracted_messages = json.loads(res.choices[0].message.content)  # type: ignore
+        logger.debug(f"Abstracted messages: {extracted_messages}")
+
+        session.events = []
+        for message in extracted_messages:
+            session.events.append(
+                Event(
+                    author="memory_optimizer",
+                    content=Content(
+                        role=message["role"],
+                        parts=[
+                            Part(
+                                text=message["content"],
+                            )
+                        ],
+                    ),
+                )
+            )
+        return session

veadk/processors/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+
+"""Runtime processors for VeADK agents.
+
+This module provides the abstract base classes for runtime processors that can
+intercept and process agent execution flows.
+"""
+
+from veadk.processors.base_run_processor import BaseRunProcessor, NoOpRunProcessor
+
+__all__ = [
+    "BaseRunProcessor",
+    "NoOpRunProcessor",
+]

veadk/processors/base_run_processor.py

@@ -0,0 +1,120 @@
+# 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.
+
+"""Base run processor for intercepting and processing agent execution."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable
+
+if TYPE_CHECKING:
+    from google.genai import types
+    from veadk.runner import Runner
+
+
+class BaseRunProcessor(ABC):
+    """Abstract base class for runtime processors.
+
+    A run processor can intercept and modify the agent execution flow by wrapping
+    the event generator function. This is useful for implementing cross-cutting
+    concerns such as:
+    - Authentication flows (e.g., OAuth2)
+    - Request/response logging
+    - Error handling and retry logic
+    - Performance monitoring
+    - Custom event filtering or transformation
+
+    The processor uses a decorator pattern to wrap the event generator, allowing
+    it to:
+    1. Intercept events from runner.run_async
+    2. Process or modify events
+    3. Inject additional events (e.g., authentication requests)
+    4. Control the execution flow (e.g., retry loops)
+
+    Example:
+        class MyProcessor(BaseRunProcessor):
+            def process_run(self, runner, message, **kwargs):
+                def decorator(event_generator_func):
+                    async def wrapper():
+                        # Pre-processing
+                        async for event in event_generator_func():
+                            # Process each event
+                            yield event
+                        # Post-processing
+                    return wrapper
+                return decorator
+    """
+
+    @abstractmethod
+    def process_run(
+        self,
+        runner: Runner,
+        message: types.Content,
+        **kwargs: Any,
+    ) -> Callable[[Callable[[], AsyncGenerator]], Callable[[], AsyncGenerator]]:
+        """Process the agent run by wrapping the event generator.
+
+        This method returns a decorator that wraps the event generator function.
+        The decorator can intercept events, modify them, or inject new events.
+
+        Args:
+            runner: The Runner instance executing the agent.
+            message: The initial message to send to the agent.
+            **kwargs: Additional keyword arguments that may be needed by specific
+                     implementations (e.g., task_updater for status updates).
+
+        Returns:
+            A decorator function that takes an event generator function and returns
+            a wrapped event generator function.
+
+        Example:
+            @processor.process_run(runner=runner, message=message)
+            async def event_generator():
+                async for event in runner.run_async(...):
+                    yield event
+        """
+        pass
+
+
+class NoOpRunProcessor(BaseRunProcessor):
+    """No-op run processor that doesn't modify the event generator.
+
+    This is the default processor used when no specific processing is needed.
+    It simply passes through all events without any modification.
+    """
+
+    def process_run(
+        self,
+        runner: Runner,
+        message: types.Content,
+        **kwargs: Any,
+    ) -> Callable[[Callable[[], AsyncGenerator]], Callable[[], AsyncGenerator]]:
+        """Return a decorator that does nothing.
+
+        Args:
+            runner: The Runner instance (unused).
+            message: The initial message (unused).
+            **kwargs: Additional keyword arguments (unused).
+
+        Returns:
+            A decorator that returns the original function unchanged.
+        """
+
+        def decorator(
+            event_generator_func: Callable[[], AsyncGenerator],
+        ) -> Callable[[], AsyncGenerator]:
+            return event_generator_func
+
+        return decorator

veadk/prompts/__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/prompts/agent_default_prompt.py

@@ -0,0 +1,30 @@
+# 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.
+
+DEFAULT_INSTRUCTION = """You an AI agent created by the VeADK team.
+
+You excel at the following tasks:
+1. Data science
+- Information gathering and fact-checking
+- Data processing and analysis
+2. Documentation
+- Writing multi-chapter articles and in-depth research reports
+3. Coding & Programming
+- Creating websites, applications, and tools
+- Solve problems and bugs in code (e.g., Python, JavaScript, SQL, ...)
+- If necessary, using programming to solve various problems beyond development
+4. If user gives you tools, finish various tasks that can be accomplished using tools and available resources
+"""
+
+DEFAULT_DESCRIPTION = """An AI agent developed by the VeADK team, specialized in data science, documentation, and software development."""

veadk/prompts/prompt_evaluator.py

@@ -0,0 +1,20 @@
+# 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.
+
+eval_principle_prompt = """
+You are a LLM for evaluating other models' responses. Note:
+- The response maybe generated by some uncertainty tools (e.g., online-search, random number), you just need to consider whether the response is human-readable, rather than focus on the specific content. Because the specific content maybe different at different time.
+"""
+
+criteria_prompt = "Determine whether the actual output is factually correct based on the expected output."

veadk/prompts/prompt_memory_processor.py

@@ -0,0 +1,55 @@
+# 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 jinja2 import Template
+
+MEMORY_PROCESSOR_SYSTEM_PROMPT = """I will give you a series of messages of memory, including messages from user and assistant.
+
+You should help me to recognize important information from the messages, and build some new messages.
+
+For example, for the following messages:
+[
+    {
+        "role": "user",
+        "content": "Hello, tell me the weather of Beijing, and remember my secret is `abc001`"
+    },
+    {
+        "role": "assistant",
+        "content": "The weather of Beijing is sunny, and the temperature is 25 degree Celsius. I have remember that your secret is `abc001`."
+    }
+]
+
+You should extract the important information from the messages, and build new messages if needed (in JSON format):
+[
+    {
+        "role": "user",
+        "content": "My secret is `abc001`."
+    }
+]
+
+The actual messages are:
+{{ messages }}
+"""
+
+
+def render_prompt(messages: list[dict]):
+    template = Template(MEMORY_PROCESSOR_SYSTEM_PROMPT)
+
+    context = {
+        "messages": messages,
+    }
+
+    rendered_prompt = template.render(context)
+
+    return rendered_prompt

veadk/prompts/prompt_optimization.py

@@ -0,0 +1,150 @@
+# 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.
+
+"""
+This file is used to optimize prompt in AgentPilot.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Callable
+
+from google.adk.tools import FunctionTool
+from google.adk.tools.mcp_tool import MCPToolset
+from jinja2 import Template
+
+from veadk import Agent
+
+# prompt = """
+# <role>
+# An experienced prompt optimizer.
+# </role>
+
+# <task>
+# Please optimize prompt to make it more efficient.
+# The prompt will be used as a system prompt and instruction of an agent.
+# The definition and context (i.e., tools) of the agent will be provided.
+# </task>
+
+# <agent_info>
+# name: {{ agent.name }}
+# model: {{ agent.model }}
+# description: {{ agent.description }}
+# </agent_info>
+
+# <agent_tools_info>
+# {% for tool in tools %}
+# <tool>
+# name: {{ tool.name }}
+# type: {{ tool.type }}
+# description: {{ tool.description }}
+# arguments: {{ tool.arguments }}
+# </tool>
+# {% endfor %}
+# </agent_tools_info>
+# """.strip()
+
+prompt = """
+Please help me to optimize the following agent prompt:
+{{ original_prompt }}
+
+
+The following information is your references：
+<agent_info>
+name: {{ agent.name }}
+model: {{ agent.model }}
+description: {{ agent.description }}
+</agent_info>
+
+<agent_tools_info>
+{% for tool in tools %}
+<tool>
+name: {{ tool.name }}
+type: {{ tool.type }}
+description: {{ tool.description }}
+arguments: {{ tool.arguments }}
+</tool>
+{% endfor %}
+</agent_tools_info>
+
+Please note that in your optimized prompt:
+- the above referenced information is not necessary. For example, the tools list of agent is not necessary in the optimized prompt, because it maybe too long. You should use the tool information to optimize the original prompt rather than simply add tool list in prompt.
+- The max length of optimized prompt should be less 4096 tokens.
+""".strip()
+
+prompt_with_feedback = """
+After you optimization, my current prompt is:
+{{ prompt }}
+
+I did some evaluations with the optimized prompt, and the feedback is: {{ feedback }}
+
+Please continue to optimize the prompt based on the feedback.
+""".strip()
+
+
+def render_prompt_feedback_with_jinja2(agent: Agent, feedback: str):
+    template = Template(prompt_with_feedback)
+
+    context = {
+        "prompt": agent.instruction,
+        "feedback": feedback,
+    }
+
+    rendered_prompt = template.render(context)
+
+    return rendered_prompt
+
+
+def render_prompt_with_jinja2(agent: Agent):
+    template = Template(prompt)
+
+    tools = []
+    for tool in agent.tools:
+        _tool_type = ""
+        _tools = []
+        if isinstance(tool, Callable):
+            _tool_type = "function"
+            _tools = [FunctionTool(tool)]
+
+        elif isinstance(tool, MCPToolset):
+            _tool_type = "tool"
+            _tools = asyncio.run(tool.get_tools())
+
+        for _tool in _tools:
+            if _tool and _tool._get_declaration():
+                tools.append(
+                    {
+                        "name": _tool.name,
+                        "description": _tool.description,
+                        "arguments": str(
+                            _tool._get_declaration().model_dump()["parameters"]  # type: ignore
+                        ),
+                        "type": _tool_type,
+                    }
+                )
+
+    context = {
+        "original_prompt": agent.instruction,
+        "agent": {
+            "name": agent.name,
+            "model": agent.model_name,
+            "description": agent.description,
+        },
+        "tools": tools,
+    }
+
+    rendered_prompt = template.render(context)
+
+    return rendered_prompt