veadk-python 0.1.0__py3-none-any.whl

veadk/memory/short_term_memory.py

@@ -0,0 +1,124 @@
+# 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 os
+from typing import Literal
+
+from google.adk.sessions import DatabaseSessionService, InMemorySessionService
+
+from veadk.config import getenv
+from veadk.utils.logger import get_logger
+
+from .short_term_memory_processor import ShortTermMemoryProcessor
+
+logger = get_logger(__name__)
+
+DEFAULT_LOCAL_DATABASE_PATH = "/tmp/veadk_local_database.db"
+
+
+class ShortTermMemory:
+    """
+    Short term memory class.
+
+    This class is used to store short term memory.
+    """
+
+    def __init__(
+        self,
+        backend: Literal["local", "database", "mysql"] = "local",
+        db_url: str = "",
+        enable_memory_optimization: bool = False,
+    ):
+        self.backend = backend
+        self.db_url = db_url
+
+        if self.backend == "mysql":
+            host = getenv("DATABASE_MYSQL_HOST")
+            user = getenv("DATABASE_MYSQL_USER")
+            password = getenv("DATABASE_MYSQL_PASSWORD")
+            database = getenv("DATABASE_MYSQL_DATABASE")
+            db_url = f"mysql+pymysql://{user}:{password}@{host}/{database}"
+
+            self.db_url = db_url
+            self.backend = "database"
+
+        if self.backend == "local":
+            logger.warning(
+                f"Short term memory backend: {self.backend}, the history will be lost after application shutdown."
+            )
+            self.session_service = InMemorySessionService()
+        elif self.backend == "database":
+            if self.db_url == "" or self.db_url is None:
+                logger.warning(
+                    f"The `db_url` {self.db_url} is an empty or None string."
+                )
+                self._use_default_database()
+            else:
+                try:
+                    self.session_service = DatabaseSessionService(db_url=self.db_url)
+                    logger.info("Connected to database.")
+                except Exception as e:
+                    logger.error(
+                        f"Failed to connect to database {self.db_url}, error: {e}."
+                    )
+                    self._use_default_database()
+        else:
+            raise ValueError(f"Unknown short term memory backend: {self.backend}")
+
+        if enable_memory_optimization and backend == "database":
+            self.processor = ShortTermMemoryProcessor()
+            intercept_get_session = self.processor.patch()
+            self.session_service.get_session = intercept_get_session(
+                self.session_service.get_session
+            )
+
+    def _use_default_database(self):
+        self.db_url = DEFAULT_LOCAL_DATABASE_PATH
+        logger.info(f"Using default local database {self.db_url}")
+        if not os.path.exists(self.db_url):
+            self.create_local_sqlite3_db(self.db_url)
+        self.session_service = DatabaseSessionService(db_url="sqlite:///" + self.db_url)
+
+    def create_local_sqlite3_db(self, path: str):
+        import sqlite3
+
+        conn = sqlite3.connect(path)
+        conn.close()
+        logger.debug(f"Create local sqlite3 database {path} done.")
+
+    async def create_session(
+        self,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+    ):
+        if isinstance(self.session_service, DatabaseSessionService):
+            list_sessions_response = await self.session_service.list_sessions(
+                app_name=app_name, user_id=user_id
+            )
+
+            logger.debug(
+                f"Loaded {len(list_sessions_response.sessions)} sessions from db {self.db_url}."
+            )
+
+        if (
+            await self.session_service.get_session(
+                app_name=app_name, user_id=user_id, session_id=session_id
+            )
+            is None
+        ):
+            # create a new session for this running
+            await self.session_service.create_session(
+                app_name=app_name, user_id=user_id, session_id=session_id
+            )

veadk/memory/short_term_memory_processor.py

@@ -0,0 +1,90 @@
+# 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 getenv
+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
+            message = {
+                "role": content.role,
+                "content": content.parts[0].text,
+            }
+            messages.append(message)
+
+        prompt = render_prompt(messages=str(messages))
+        res = completion(
+            model=getenv("MODEL_AGENT_PROVIDER") + "/" + getenv("MODEL_AGENT_NAME"),
+            base_url=getenv("MODEL_AGENT_API_BASE"),
+            api_key=getenv("MODEL_AGENT_API_KEY"),
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt,
+                }
+            ],
+        )
+        extracted_messages = json.loads(res.choices[0].message.content)
+        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/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,158 @@
+# 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:
+        if isinstance(tool, Callable):
+            type = "function"
+
+            _tool = FunctionTool(tool)
+            tools.append(
+                {
+                    "name": _tool.name,
+                    "description": _tool.description,
+                    "arguments": str(
+                        _tool._get_declaration().model_dump()["parameters"]
+                    ),
+                    "type": type,
+                }
+            )
+        elif isinstance(tool, MCPToolset):
+            type = "tool"
+
+            _tools = asyncio.run(tool.get_tools())
+            for _tool in _tools:
+                tools.append(
+                    {
+                        "name": _tool.name,
+                        "description": _tool.description,
+                        "arguments": str(
+                            _tool._get_declaration().model_dump()["parameters"]
+                        ),
+                        "type": type,
+                    }
+                )
+
+    context = {
+        "original_prompt": agent.instruction,
+        "agent": {
+            "name": agent.name,
+            "model": agent.model_name,
+            "description": agent.description,
+            # "tools": str(agent.tools),
+        },
+        "tools": tools,
+    }
+
+    rendered_prompt = template.render(context)
+
+    return rendered_prompt