veadk-python 0.5.25__tar.gz → 0.5.26__tar.gz

{veadk_python-0.5.25/veadk_python.egg-info → veadk_python-0.5.26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: veadk-python
-Version: 0.5.25
+Version: 0.5.26
 Summary: Volcengine agent development kit, integrations with Volcengine cloud services.
 Author-email: Yaozheng Fang <fangyozheng@gmail.com>, Guodong Li <cu.eric.lee@gmail.com>, Zhi Han <sliverydayday@gmail.com>, Meng Wang <mengwangwm@gmail.com>
 License:                                  Apache License

{veadk_python-0.5.25 → veadk_python-0.5.26}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "veadk-python"
-version = "0.5.25"
+version = "0.5.26"
 description = "Volcengine agent development kit, integrations with Volcengine cloud services."
 readme = "README.md"
 requires-python = ">=3.10"

veadk_python-0.5.26/veadk/tools/vanna_tools/agent_memory.py

@@ -0,0 +1,339 @@
+# 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, Dict, Optional, List
+from google.adk.tools import BaseTool, ToolContext
+from google.genai import types
+from vanna.tools.agent_memory import (
+    SaveQuestionToolArgsTool as VannaSaveQuestionToolArgsTool,
+    SearchSavedCorrectToolUsesTool as VannaSearchSavedCorrectToolUsesTool,
+    SaveTextMemoryTool as VannaSaveTextMemoryTool,
+)
+from vanna.core.user import User
+from vanna.core.tool import ToolContext as VannaToolContext
+
+
+class SaveQuestionToolArgsTool(BaseTool):
+    """Save successful question-tool-argument combinations for future reference."""
+
+    def __init__(
+        self,
+        agent_memory,
+        access_groups: Optional[List[str]] = None,
+    ):
+        """
+        Initialize the save tool usage tool with custom agent_memory.
+
+        Args:
+            agent_memory: A Vanna agent memory instance (e.g., DemoAgentMemory)
+            access_groups: List of user groups that can access this tool (e.g., ['admin'])
+        """
+        self.agent_memory = agent_memory
+        self.vanna_tool = VannaSaveQuestionToolArgsTool()
+        self.access_groups = access_groups or ["admin", "user"]
+
+        super().__init__(
+            name="save_question_tool_args",  # Keep the same name as Vanna
+            description="Save a successful question-tool-argument combination for future reference.",
+        )
+
+    def _get_declaration(self) -> types.FunctionDeclaration:
+        return types.FunctionDeclaration(
+            name=self.name,
+            description=self.description,
+            parameters=types.Schema(
+                type=types.Type.OBJECT,
+                properties={
+                    "question": types.Schema(
+                        type=types.Type.STRING,
+                        description="The original question that was asked",
+                    ),
+                    "tool_name": types.Schema(
+                        type=types.Type.STRING,
+                        description="The name of the tool that was used successfully",
+                    ),
+                    "args": types.Schema(
+                        type=types.Type.OBJECT,
+                        description="The arguments that were passed to the tool",
+                    ),
+                },
+                required=["question", "tool_name", "args"],
+            ),
+        )
+
+    def _get_user_groups(self, tool_context: ToolContext) -> List[str]:
+        """Get user groups from context."""
+        user_groups = tool_context.state.get("user_groups", ["user"])
+        return user_groups
+
+    def _check_access(self, user_groups: List[str]) -> bool:
+        """Check if user has access to this tool."""
+        return any(group in self.access_groups for group in user_groups)
+
+    def _create_vanna_context(
+        self, tool_context: ToolContext, user_groups: List[str]
+    ) -> VannaToolContext:
+        """Create Vanna context from Veadk ToolContext."""
+        user_id = tool_context.user_id
+        session_id = tool_context.session.id
+        user_email = tool_context.state.get("user_email", "user@example.com")
+
+        vanna_user = User(
+            id=user_id + "_" + session_id,
+            email=user_email,
+            group_memberships=user_groups,
+        )
+
+        vanna_context = VannaToolContext(
+            user=vanna_user,
+            conversation_id=session_id,
+            request_id=session_id,
+            agent_memory=self.agent_memory,
+        )
+
+        return vanna_context
+
+    async def run_async(
+        self, *, args: Dict[str, Any], tool_context: ToolContext
+    ) -> str:
+        """Save a tool usage pattern."""
+        question = args.get("question", "").strip()
+        tool_name = args.get("tool_name", "").strip()
+        tool_args = args.get("args", {})
+
+        if not question:
+            return "Error: No question provided"
+
+        if not tool_name:
+            return "Error: No tool name provided"
+
+        try:
+            user_groups = self._get_user_groups(tool_context)
+
+            if not self._check_access(user_groups):
+                return f"Error: Access denied. This tool requires one of the following groups: {', '.join(self.access_groups)}"
+
+            vanna_context = self._create_vanna_context(tool_context, user_groups)
+
+            args_model = self.vanna_tool.get_args_schema()(
+                question=question, tool_name=tool_name, args=tool_args
+            )
+            result = await self.vanna_tool.execute(vanna_context, args_model)
+
+            return str(result.result_for_llm)
+        except Exception as e:
+            return f"Error saving tool usage: {str(e)}"
+
+
+class SearchSavedCorrectToolUsesTool(BaseTool):
+    """Search for similar tool usage patterns based on a question."""
+
+    def __init__(
+        self,
+        agent_memory,
+        access_groups: Optional[List[str]] = None,
+    ):
+        """
+        Initialize the search similar tools tool with custom agent_memory.
+
+        Args:
+            agent_memory: A Vanna agent memory instance (e.g., DemoAgentMemory)
+            access_groups: List of user groups that can access this tool (e.g., ['admin', 'user'])
+            user_group_resolver: Optional callable that takes ToolContext and returns user groups
+        """
+        self.agent_memory = agent_memory
+        self.vanna_tool = VannaSearchSavedCorrectToolUsesTool()
+        self.access_groups = access_groups or ["admin", "user"]
+
+        super().__init__(
+            name="search_saved_correct_tool_uses",  # Keep the same name as Vanna
+            description="Search for similar tool usage patterns based on a question.",
+        )
+
+    def _get_declaration(self) -> types.FunctionDeclaration:
+        return types.FunctionDeclaration(
+            name=self.name,
+            description=self.description,
+            parameters=types.Schema(
+                type=types.Type.OBJECT,
+                properties={
+                    "question": types.Schema(
+                        type=types.Type.STRING,
+                        description="The question to find similar tool usage patterns for",
+                    ),
+                    "limit": types.Schema(
+                        type=types.Type.INTEGER,
+                        description="Maximum number of results to return (default: 10)",
+                    ),
+                },
+                required=["question"],
+            ),
+        )
+
+    def _get_user_groups(self, tool_context: ToolContext) -> List[str]:
+        """Get user groups from context."""
+        user_groups = tool_context.state.get("user_groups", ["user"])
+        return user_groups
+
+    def _check_access(self, user_groups: List[str]) -> bool:
+        """Check if user has access to this tool."""
+        return any(group in self.access_groups for group in user_groups)
+
+    def _create_vanna_context(
+        self, tool_context: ToolContext, user_groups: List[str]
+    ) -> VannaToolContext:
+        """Create Vanna context from Veadk ToolContext."""
+        user_id = tool_context.user_id
+        session_id = tool_context.session.id
+        user_email = tool_context.state.get("user_email", "user@example.com")
+
+        vanna_user = User(
+            id=user_id + "_" + session_id,
+            email=user_email,
+            group_memberships=user_groups,
+        )
+
+        vanna_context = VannaToolContext(
+            user=vanna_user,
+            conversation_id=session_id,
+            request_id=session_id,
+            agent_memory=self.agent_memory,
+        )
+
+        return vanna_context
+
+    async def run_async(
+        self, *, args: Dict[str, Any], tool_context: ToolContext
+    ) -> str:
+        """Search for similar tool usage patterns."""
+        question = args.get("question", "").strip()
+        limit = args.get("limit", 10)
+
+        if not question:
+            return "Error: No question provided"
+
+        try:
+            user_groups = self._get_user_groups(tool_context)
+
+            if not self._check_access(user_groups):
+                return f"Error: Access denied. This tool requires one of the following groups: {', '.join(self.access_groups)}"
+
+            vanna_context = self._create_vanna_context(tool_context, user_groups)
+
+            args_model = self.vanna_tool.get_args_schema()(
+                question=question, limit=limit
+            )
+            result = await self.vanna_tool.execute(vanna_context, args_model)
+
+            return str(result.result_for_llm)
+        except Exception as e:
+            return f"Error searching similar tools: {str(e)}"
+
+
+class SaveTextMemoryTool(BaseTool):
+    """Save free-form text memories for important insights, observations, or context."""
+
+    def __init__(
+        self,
+        agent_memory,
+        access_groups: Optional[List[str]] = None,
+    ):
+        """
+        Initialize the save text memory tool with custom agent_memory.
+
+        Args:
+            agent_memory: A Vanna agent memory instance (e.g., DemoAgentMemory)
+            access_groups: List of user groups that can access this tool (e.g., ['admin', 'user'])
+            user_group_resolver: Optional callable that takes ToolContext and returns user groups
+        """
+        self.agent_memory = agent_memory
+        self.vanna_tool = VannaSaveTextMemoryTool()
+        self.access_groups = access_groups or ["admin", "user"]
+
+        super().__init__(
+            name="save_text_memory",  # Keep the same name as Vanna
+            description="Save free-form text memory for important insights, observations, or context.",
+        )
+
+    def _get_declaration(self) -> types.FunctionDeclaration:
+        return types.FunctionDeclaration(
+            name=self.name,
+            description=self.description,
+            parameters=types.Schema(
+                type=types.Type.OBJECT,
+                properties={
+                    "content": types.Schema(
+                        type=types.Type.STRING,
+                        description="The text content to save as a memory",
+                    ),
+                },
+                required=["content"],
+            ),
+        )
+
+    def _get_user_groups(self, tool_context: ToolContext) -> List[str]:
+        """Get user groups from context."""
+        user_groups = tool_context.state.get("user_groups", ["user"])
+        return user_groups
+
+    def _check_access(self, user_groups: List[str]) -> bool:
+        """Check if user has access to this tool."""
+        return any(group in self.access_groups for group in user_groups)
+
+    def _create_vanna_context(
+        self, tool_context: ToolContext, user_groups: List[str]
+    ) -> VannaToolContext:
+        """Create Vanna context from Veadk ToolContext."""
+        user_id = tool_context.user_id
+        session_id = tool_context.session.id
+        user_email = tool_context.state.get("user_email", "user@example.com")
+
+        vanna_user = User(
+            id=user_id + "_" + session_id,
+            email=user_email,
+            group_memberships=user_groups,
+        )
+
+        vanna_context = VannaToolContext(
+            user=vanna_user,
+            conversation_id=session_id,
+            request_id=session_id,
+            agent_memory=self.agent_memory,
+        )
+
+        return vanna_context
+
+    async def run_async(
+        self, *, args: Dict[str, Any], tool_context: ToolContext
+    ) -> str:
+        """Save a text memory."""
+        content = args.get("content", "").strip()
+
+        if not content:
+            return "Error: No content provided"
+
+        try:
+            user_groups = self._get_user_groups(tool_context)
+
+            if not self._check_access(user_groups):
+                return f"Error: Access denied. This tool requires one of the following groups: {', '.join(self.access_groups)}"
+
+            vanna_context = self._create_vanna_context(tool_context, user_groups)
+
+            args_model = self.vanna_tool.get_args_schema()(content=content)
+            result = await self.vanna_tool.execute(vanna_context, args_model)
+
+            return str(result.result_for_llm)
+        except Exception as e:
+            return f"Error saving text memory: {str(e)}"

veadk_python-0.5.26/veadk/tools/vanna_tools/examples/example.py

@@ -0,0 +1,228 @@
+# 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 veadk import Agent, Runner
+from veadk.tools.vanna_tools.vanna_toolset import VannaToolSet
+from google.adk.sessions import InMemorySessionService
+
+
+# Create a session with user groups for access control
+async def create_session(user_groups: list = ["user"]):
+    session_service = InMemorySessionService()
+    example_session = await session_service.create_session(
+        app_name="example_app",
+        user_id="example_user",
+        state={"user_groups": user_groups},
+    )
+    return session_service, example_session
+
+
+vanna_toolset = VannaToolSet(
+    connection_string="sqlite:///tmp/Chinook.sqlite", file_storage="/tmp/vanna_files"
+)
+
+
+# Define the Veadk Agent with class-based tools
+agent: Agent = Agent(
+    name="vanna_sql_agent",
+    description="An intelligent agent that can query databases, visualize data, and generate reports.",
+    instruction="""
+    You are a helpful assistant that can answer questions about data in the Chinook database.
+    You can execute SQL queries, visualize the results, save/search useful tool usage patterns, and generate documents.
+
+    Here is the schema of the Chinook database:
+    ```sql
+    CREATE TABLE [Album]
+    (
+        [AlbumId] INTEGER  NOT NULL,
+        [Title] NVARCHAR(160)  NOT NULL,
+        [ArtistId] INTEGER  NOT NULL,
+        CONSTRAINT [PK_Album] PRIMARY KEY  ([AlbumId]),
+        FOREIGN KEY ([ArtistId]) REFERENCES [Artist] ([ArtistId]) 
+                    ON DELETE NO ACTION ON UPDATE NO ACTION
+    );
+    CREATE TABLE [Artist]
+    (
+        [ArtistId] INTEGER  NOT NULL,
+        [Name] NVARCHAR(120),
+        CONSTRAINT [PK_Artist] PRIMARY KEY  ([ArtistId])
+    );
+    CREATE TABLE [Customer]
+    (
+        [CustomerId] INTEGER  NOT NULL,
+        [FirstName] NVARCHAR(40)  NOT NULL,
+        [LastName] NVARCHAR(20)  NOT NULL,
+        [Company] NVARCHAR(80),
+        [Address] NVARCHAR(70),
+        [City] NVARCHAR(40),
+        [State] NVARCHAR(40),
+        [Country] NVARCHAR(40),
+        [PostalCode] NVARCHAR(10),
+        [Phone] NVARCHAR(24),
+        [Fax] NVARCHAR(24),
+        [Email] NVARCHAR(60)  NOT NULL,
+        [SupportRepId] INTEGER,
+        CONSTRAINT [PK_Customer] PRIMARY KEY  ([CustomerId]),
+        FOREIGN KEY ([SupportRepId]) REFERENCES [Employee] ([EmployeeId]) 
+                    ON DELETE NO ACTION ON UPDATE NO ACTION
+    );
+    CREATE TABLE [Employee]
+    (
+        [EmployeeId] INTEGER  NOT NULL,
+        [LastName] NVARCHAR(20)  NOT NULL,
+        [FirstName] NVARCHAR(20)  NOT NULL,
+        [Title] NVARCHAR(30),
+        [ReportsTo] INTEGER,
+        [BirthDate] DATETIME,
+        [HireDate] DATETIME,
+        [Address] NVARCHAR(70),
+        [City] NVARCHAR(40),
+        [State] NVARCHAR(40),
+        [Country] NVARCHAR(40),
+        [PostalCode] NVARCHAR(10),
+        [Phone] NVARCHAR(24),
+        [Fax] NVARCHAR(24),
+        [Email] NVARCHAR(60),
+        CONSTRAINT [PK_Employee] PRIMARY KEY  ([EmployeeId]),
+        FOREIGN KEY ([ReportsTo]) REFERENCES [Employee] ([EmployeeId]) 
+                    ON DELETE NO ACTION ON UPDATE NO ACTION
+    );
+    CREATE TABLE [Genre]
+    (
+        [GenreId] INTEGER  NOT NULL,
+        [Name] NVARCHAR(120),
+        CONSTRAINT [PK_Genre] PRIMARY KEY  ([GenreId])
+    );
+    CREATE TABLE [Invoice]
+    (
+        [InvoiceId] INTEGER  NOT NULL,
+        [CustomerId] INTEGER  NOT NULL,
+        [InvoiceDate] DATETIME  NOT NULL,
+        [BillingAddress] NVARCHAR(70),
+        [BillingCity] NVARCHAR(40),
+        [BillingState] NVARCHAR(40),
+        [BillingCountry] NVARCHAR(40),
+        [BillingPostalCode] NVARCHAR(10),
+        [Total] NUMERIC(10,2)  NOT NULL,
+        CONSTRAINT [PK_Invoice] PRIMARY KEY  ([InvoiceId]),
+        FOREIGN KEY ([CustomerId]) REFERENCES [Customer] ([CustomerId]) 
+                    ON DELETE NO ACTION ON UPDATE NO ACTION
+    );
+    CREATE TABLE [InvoiceLine]
+    (
+        [InvoiceLineId] INTEGER  NOT NULL,
+        [InvoiceId] INTEGER  NOT NULL,
+        [TrackId] INTEGER  NOT NULL,
+        [UnitPrice] NUMERIC(10,2)  NOT NULL,
+        [Quantity] INTEGER  NOT NULL,
+        CONSTRAINT [PK_InvoiceLine] PRIMARY KEY  ([InvoiceLineId]),
+        FOREIGN KEY ([InvoiceId]) REFERENCES [Invoice] ([InvoiceId]) 
+                    ON DELETE NO ACTION ON UPDATE NO ACTION,
+        FOREIGN KEY ([TrackId]) REFERENCES [Track] ([TrackId]) 
+                    ON DELETE NO ACTION ON UPDATE NO ACTION
+    );
+    CREATE TABLE [MediaType]
+    (
+        [MediaTypeId] INTEGER  NOT NULL,
+        [Name] NVARCHAR(120),
+        CONSTRAINT [PK_MediaType] PRIMARY KEY  ([MediaTypeId])
+    );
+    CREATE TABLE [Playlist]
+    (
+        [PlaylistId] INTEGER  NOT NULL,
+        [Name] NVARCHAR(120),
+        CONSTRAINT [PK_Playlist] PRIMARY KEY  ([PlaylistId])
+    );
+    CREATE TABLE [PlaylistTrack]
+    (
+        [PlaylistId] INTEGER  NOT NULL,
+        [TrackId] INTEGER  NOT NULL,
+        CONSTRAINT [PK_PlaylistTrack] PRIMARY KEY  ([PlaylistId], [TrackId]),
+        FOREIGN KEY ([PlaylistId]) REFERENCES [Playlist] ([PlaylistId]) 
+                    ON DELETE NO ACTION ON UPDATE NO ACTION,
+        FOREIGN KEY ([TrackId]) REFERENCES [Track] ([TrackId]) 
+                    ON DELETE NO ACTION ON UPDATE NO ACTION
+    );
+    CREATE TABLE [Track]
+    (
+        [TrackId] INTEGER  NOT NULL,
+        [Name] NVARCHAR(200)  NOT NULL,
+        [AlbumId] INTEGER,
+        [MediaTypeId] INTEGER  NOT NULL,
+        [GenreId] INTEGER,
+        [Composer] NVARCHAR(220),
+        [Milliseconds] INTEGER  NOT NULL,
+        [Bytes] INTEGER,
+        [UnitPrice] NUMERIC(10,2)  NOT NULL,
+        CONSTRAINT [PK_Track] PRIMARY KEY  ([TrackId]),
+        FOREIGN KEY ([AlbumId]) REFERENCES [Album] ([AlbumId]) 
+                    ON DELETE NO ACTION ON UPDATE NO ACTION,
+        FOREIGN KEY ([GenreId]) REFERENCES [Genre] ([GenreId]) 
+                    ON DELETE NO ACTION ON UPDATE NO ACTION,
+        FOREIGN KEY ([MediaTypeId]) REFERENCES [MediaType] ([MediaTypeId]) 
+                    ON DELETE NO ACTION ON UPDATE NO ACTION
+    );
+    ```
+
+    Here are some examples of how to query this database:
+    
+    Q: Get all the tracks in the album 'Balls to the Wall'.
+    A: SELECT * FROM Track WHERE AlbumId = (SELECT AlbumId FROM Album WHERE Title = 'Balls to the Wall')
+
+    Q: Get the total sales for each customer.
+    A: SELECT c.FirstName, c.LastName, SUM(i.Total) as TotalSales FROM Customer c JOIN Invoice i ON c.CustomerId = i.CustomerId GROUP BY c.CustomerId
+
+    Q: How many tracks are there in each genre?
+    A: SELECT g.Name, COUNT(t.TrackId) as TrackCount FROM Genre g JOIN Track t ON g.GenreId = t.GenreId GROUP BY g.GenreId
+    
+    Available tools (using Vanna's original names):
+    1. `run_sql` - Execute SQL queries
+    2. `visualize_data` - Create visualizations from CSV files
+    3. `write_file` - Save content to files
+    4. `save_question_tool_args` - Save successful tool usage patterns (admin only)
+    5. `search_saved_correct_tool_uses` - Search for similar tool usage patterns
+    6. `summarize_data` - Generate statistical summaries of CSV files
+    """,
+    tools=[vanna_toolset],
+    model_extra_config={"extra_body": {"thinking": {"type": "disabled"}}},
+)
+
+
+async def main(prompt: str, user_groups: list = None) -> str:
+    session_service, example_session = await create_session(
+        user_groups
+    )  # Default to 'user' group if not specified
+
+    runner = Runner(
+        agent=agent,
+        app_name=example_session.app_name,
+        user_id=example_session.user_id,
+        session_service=session_service,
+    )
+
+    response = await runner.run(
+        messages=prompt,
+        session_id=example_session.id,
+    )
+
+    return response
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    print("=== Example 1: Regular User ===")
+    user_input = "What are the top 5 selling albums?"
+    response = asyncio.run(main(user_input, user_groups=["user"]))
+    print(response)