rebase knowledgebase tool

Author: yaozheng-fang

veadk/agent.py

@@ -133,10 +133,14 @@ def model_post_init(self, __context: Any) -> None:
             )
 
         if self.knowledgebase:
-            from veadk.tools import load_knowledgebase_tool
+            from veadk.tools.builtin_tools.load_knowledgebase import (
+                LoadKnowledgebaseTool,
+            )
 
-            load_knowledgebase_tool.knowledgebase = self.knowledgebase
-            self.tools.append(load_knowledgebase_tool.load_knowledgebase_tool)
+            load_knowledgebase_tool = LoadKnowledgebaseTool(
+                knowledgebase=self.knowledgebase
+            )
+            self.tools.append(load_knowledgebase_tool)
 
         if self.long_term_memory is not None:
             from google.adk.tools import load_memory

veadk/knowledgebase/knowledgebase.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Callable, Literal
+from __future__ import annotations
+
+from typing import Any, Callable, Literal, Union
 
 from pydantic import BaseModel, Field
-from typing_extensions import Union
 
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
 from veadk.knowledgebase.entry import KnowledgebaseEntry
@@ -54,11 +55,11 @@ def _get_backend_cls(backend: str) -> type[BaseKnowledgebaseBackend]:
     raise ValueError(f"Unsupported knowledgebase backend: {backend}")
 
 
-def build_knowledgebase_index(app_name: str):
-    return f"veadk_kb_{app_name}"
+class KnowledgeBase(BaseModel):
+    name: str = "user_knowledgebase"
 
+    description: str = "This knowledgebase stores some user-related information."
 
-class KnowledgeBase(BaseModel):
     backend: Union[
         Literal["local", "opensearch", "viking", "redis"], BaseKnowledgebaseBackend
     ] = "local"

veadk/tools/builtin_tools/load_knowledgebase.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+from google.adk.models.llm_request import LlmRequest
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+from pydantic import BaseModel, Field
+from typing_extensions import override
+
+from veadk.knowledgebase import KnowledgeBase
+from veadk.knowledgebase.entry import KnowledgebaseEntry
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class LoadKnowledgebaseResponse(BaseModel):
+    knowledges: list[KnowledgebaseEntry] = Field(default_factory=list)
+
+
+class LoadKnowledgebaseTool(FunctionTool):
+    """A tool that loads the common knowledgebase"""
+
+    def __init__(self, knowledgebase: KnowledgeBase):
+        super().__init__(self.load_knowledgebase)
+
+        self.knowledgebase = knowledgebase
+
+        if not self.custom_metadata:
+            self.custom_metadata = {}
+        self.custom_metadata["backend"] = knowledgebase.backend
+
+    @override
+    def _get_declaration(self) -> types.FunctionDeclaration | None:
+        return types.FunctionDeclaration(
+            name=self.name,
+            description=self.description,
+            parameters=types.Schema(
+                type=types.Type.OBJECT,
+                properties={
+                    "query": types.Schema(
+                        type=types.Type.STRING,
+                    )
+                },
+                required=["query"],
+            ),
+        )
+
+    @override
+    async def process_llm_request(
+        self,
+        *,
+        tool_context: ToolContext,
+        llm_request: LlmRequest,
+    ) -> None:
+        await super().process_llm_request(
+            tool_context=tool_context, llm_request=llm_request
+        )
+        # Tell the model about the knowledgebase.
+        llm_request.append_instructions(
+            [
+                f"""
+You have a knowledgebase (knowledegebase name is `{self.knowledgebase.name}`, knowledgebase description is `{self.knowledgebase.description}`). You can use it to answer questions. If any questions need
+you to look up the knowledgebase, you should call load_knowledgebase function with a query.
+"""
+            ]
+        )
+
+    async def load_knowledgebase(
+        self, query: str, tool_context: ToolContext
+    ) -> LoadKnowledgebaseResponse:
+        """Loads the knowledgebase for the user.
+
+        Args:
+        query: The query to load the knowledgebase for.
+
+        Returns:
+        A list of knowledgebase results.
+        """
+        logger.info(f"Search knowledgebase: {self.knowledgebase.name}")
+        response = self.knowledgebase.search(query)
+        logger.info(f"Loaded {len(response)} knowledgebase entries for query: {query}")
+        return LoadKnowledgebaseResponse(knowledges=response)