load_knowledgebase.py

# 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

import json

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.tools.builtin_tools.load_kb_queries import load_profile
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
        )

        index = self.knowledgebase.index
        if self.knowledgebase.enable_profile:
            from pathlib import Path

            profile_names = []
            profile_descriptions = []

            with open(
                f"./profiles/knowledgebase/profiles_{index}/profile_list.json",
                "r",
            ) as f:
                profile_names = json.load(f)

            for profile_name in profile_names:
                profile_descriptions.append(
                    load_profile(
                        Path(
                            f"./profiles/knowledgebase/profiles_{index}/profile_{profile_name}.json"
                        ),
                    )["description"]
                )

            profiles_text = "\n".join(
                f"- profile_name: {name}\n  profile_description: {profile_descriptions[idx]}"
                for idx, name in enumerate(profile_names)
            )

        # 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.
"""
            ],
        )

        if self.knowledgebase.enable_profile:
            llm_request.append_instructions(
                [
                    f"""
The knowledgebase is divided into the following profiles: 

{profiles_text}

You should choose some profiles which are relevant to the user question. Before load the knowledgebase, you must call `load_kb_queries` to load the recommanded queries of the knowledgebase profiles. You should generate final knowledgebase queries based on the user question and recommanded queries.
"""
                ]
            )

        if self.knowledgebase.query_with_user_profile:
            from veadk import Agent

            agent = tool_context._invocation_context.agent
            if not isinstance(agent, Agent) or not agent.long_term_memory:
                logger.error(
                    "Agent in tool context is not an instance of veadk.Agent or long term memory is not set in agent attribution. Cannot load user profile."
                )
                return

            user_profile = agent.long_term_memory.get_user_profile(tool_context.user_id)

            if user_profile:
                llm_request.append_instructions(
                    [
                        f"""
Please generate the knowledgebase queries based on the user profile (description) at the same time. For example, for a query `quick sort algorithm`, you should generate `quick sort algorithm for python` if the user is a python developer, or `quick sort algorithm friendly introduction` if the user is a beginner.

The user profile is : 

{user_profile}
"""
                    ]
                )

    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)