veadk-python 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

veadk/knowledgebase/knowledgebase.py

@@ -14,28 +14,36 @@
 
 from typing import BinaryIO, Literal, TextIO
 
+from veadk.database.database_adapter import get_knowledgebase_database_adapter
 from veadk.database.database_factory import DatabaseFactory
 from veadk.utils.logger import get_logger
 
-from .knowledgebase_database_adapter import get_knowledgebase_adapter
-
 logger = get_logger(__name__)
 
 
+def build_knowledgebase_index(app_name: str):
+    return f"veadk_kb_{app_name}"
+
+
 class KnowledgeBase:
     def __init__(
         self,
         backend: Literal["local", "opensearch", "viking", "redis", "mysql"] = "local",
-        top_k: int = 5,
+        top_k: int = 10,
         db_config=None,
     ):
-        logger.debug(f"Create knowledgebase, backend is {backend}")
+        logger.info(
+            f"Initializing knowledgebase: backend={backend} top_k={top_k} db_config={db_config}"
+        )
+
         self.backend = backend
         self.top_k = top_k
 
         self.db_client = DatabaseFactory.create(backend=backend, config=db_config)
-        self.adapter = get_knowledgebase_adapter(backend)(
-            database_client=self.db_client
+        self.adapter = get_knowledgebase_database_adapter(self.db_client)
+
+        logger.info(
+            f"Initialized knowledgebase: db_client={self.db_client} adapter={self.adapter}"
         )
 
     def add(
@@ -46,38 +54,31 @@ class KnowledgeBase:
     ):
         """
         Add documents to the vector database.
-        You can only upload files or file characters when the adapter type you use is vikingdb.
+        You can only upload files or file characters when the adapter type used is vikingdb.
         In addition, if you upload data of the bytes type,
             for example, if you read the file stream of a pdf, then you need to pass an additional parameter file_ext = '.pdf'.
         """
-        kwargs.pop("session_id", None)  # remove session_id
-        self.adapter.add(
-            data, app_name, user_id="user_id", session_id="session_id", **kwargs
-        )
+        if self.backend != "viking" and not (
+            isinstance(data, str) or isinstance(data, list)
+        ):
+            raise ValueError(
+                "Only vikingdb supports uploading files or file characters."
+            )
 
-    def search(self, query: str, app_name: str, top_k: int = None) -> list[str]:
-        """Retrieve documents similar to the query text in the vector database.
+        index = build_knowledgebase_index(app_name)
 
-        Args:
-            query (str): The query text to be retrieved (e.g., "Who proposed the Turing machine model?")
+        logger.info(f"Adding documents to knowledgebase: index={index}")
 
-        Returns:
-            list[str]: A list of the top most similar document contents retrieved (sorted by vector similarity)
-        """
+        self.adapter.add(data=data, index=index)
+
+    def search(self, query: str, app_name: str, top_k: int = None) -> list[str]:
         top_k = self.top_k if top_k is None else top_k
 
-        result = self.adapter.query(
-            query=query, app_name=app_name, user_id="user_id", top_k=top_k
+        logger.info(
+            f"Searching knowledgebase: app_name={app_name} query={query} top_k={top_k}"
         )
+        index = build_knowledgebase_index(app_name)
+        result = self.adapter.query(query=query, index=index, top_k=top_k)
         if len(result) == 0:
             logger.warning(f"No documents found in knowledgebase. Query: {query}")
         return result
-
-    def delete(self, app_name: str, user_id: str, session_id: str):
-        """Delete documents in the vector database.
-        Args:
-            app_name (str): The name of the application
-            user_id (str): The user ID
-            session_id (str): The session ID
-        """
-        self.adapter.delete(app_name=app_name, user_id=user_id, session_id=session_id)

veadk/memory/long_term_memory.py

@@ -16,6 +16,7 @@
 import json
 from typing import Literal
 
+from google.adk.events.event import Event
 from google.adk.memory.base_memory_service import (
     BaseMemoryService,
     SearchMemoryResponse,
@@ -25,67 +26,103 @@ from google.adk.sessions import Session
 from google.genai import types
 from typing_extensions import override
 
-from veadk.config import getenv
 from veadk.database import DatabaseFactory
+from veadk.database.database_adapter import get_long_term_memory_database_adapter
 from veadk.utils.logger import get_logger
 
-from .memory_database_adapter import get_memory_adapter
-
 logger = get_logger(__name__)
 
 
+def build_long_term_memory_index(app_name: str, user_id: str):
+    return f"{app_name}_{user_id}"
+
+
 class LongTermMemory(BaseMemoryService):
     def __init__(
         self,
         backend: Literal[
             "local", "opensearch", "redis", "mysql", "viking"
         ] = "opensearch",
-        top_k: int = getenv("LONGTERM_MEMORY_TOP_K", 3),
+        top_k: int = 5,
     ):
         if backend == "viking":
             backend = "viking_mem"
         self.top_k = top_k
         self.backend = backend
 
+        logger.info(
+            f"Initializing long term memory: backend={self.backend} top_k={self.top_k}"
+        )
+
         self.db_client = DatabaseFactory.create(
             backend=backend,
         )
-        logger.info(f"Long term memory backend is `{backend}`.")
+        self.adapter = get_long_term_memory_database_adapter(self.db_client)
+
+        logger.info(
+            f"Initialized long term memory: db_client={self.db_client} adapter={self.adapter}"
+        )
+
+    def _filter_and_convert_events(self, events: list[Event]) -> list[str]:
+        final_events = []
+        for event in events:
+            # filter: bad event
+            if not event.content or not event.content.parts:
+                continue
+
+            # filter: only add user event to memory to enhance retrieve performance
+            if not event.author == "user":
+                continue
+
+            # filter: discard function call and function response
+            if not event.content.parts[0].text:
+                continue
+
+            # convert: to string-format for storage
+            message = event.content.model_dump(exclude_none=True, mode="json")
 
-        self.adapter = get_memory_adapter(backend)(database_client=self.db_client)
+            final_events.append(json.dumps(message, ensure_ascii=False))
+        return final_events
 
     @override
     async def add_session_to_memory(
         self,
         session: Session,
     ):
-        event_list = []
-        for event in session.events:
-            if not event.content or not event.content.parts:
-                continue
+        event_strings = self._filter_and_convert_events(session.events)
+        index = build_long_term_memory_index(session.app_name, session.user_id)
 
-            message = event.content.model_dump(exclude_none=True, mode="json")
-            if (
-                "text" not in message["parts"][0]
-            ):  # remove function_call & function_resp
-                continue
-            event_list.append(json.dumps(message))
-        self.adapter.add(
-            event_list,
-            app_name=session.app_name,
-            user_id=session.user_id,
-            session_id=session.id,
+        logger.info(
+            f"Adding {len(event_strings)} events to long term memory: index={index}"
+        )
+
+        # check if viking memory database, should give a user id： if/else
+        if self.backend == "viking_mem":
+            self.adapter.add(data=event_strings, index=index, user_id=session.user_id)
+        else:
+            self.adapter.add(data=event_strings, index=index)
+
+        logger.info(
+            f"Added {len(event_strings)} events to long term memory: index={index}"
         )
 
     @override
     async def search_memory(self, *, app_name: str, user_id: str, query: str):
-        memory_chunks = self.adapter.query(
-            query=query,
-            app_name=app_name,
-            user_id=user_id,
+        index = build_long_term_memory_index(app_name, user_id)
+
+        logger.info(
+            f"Searching long term memory: query={query} index={index} top_k={self.top_k}"
         )
-        if len(memory_chunks) == 0:
-            return SearchMemoryResponse()
+
+        # user id if viking memory db
+        if self.backend == "viking_mem":
+            memory_chunks = self.adapter.query(
+                query=query, index=index, top_k=self.top_k, user_id=user_id
+            )
+        else:
+            memory_chunks = self.adapter.query(
+                query=query, index=index, top_k=self.top_k
+            )
 
         memory_events = []
         for memory in memory_chunks:
@@ -94,26 +131,25 @@ class LongTermMemory(BaseMemoryService):
                 try:
                     text = memory_dict["parts"][0]["text"]
                     role = memory_dict["role"]
-                except KeyError as e:
-                    logger.error(
-                        f"Memory content: {memory_dict}. Error parsing memory: {e}"
+                except KeyError as _:
+                    # prevent not a standard text-based event
+                    logger.warning(
+                        f"Memory content: {memory_dict}. Skip return this memory."
                     )
                     continue
             except json.JSONDecodeError:
+                # prevent the memory string is not dumped by `Event` class
                 text = memory
                 role = "user"
+
             memory_events.append(
                 MemoryEntry(
                     author="user",
                     content=types.Content(parts=[types.Part(text=text)], role=role),
                 )
             )
-        return SearchMemoryResponse(memories=memory_events)
 
-    @override
-    async def delete_memory(self, *, app_name: str, user_id: str):
-        self.adapter.delete(
-            app_name=app_name,
-            user_id=user_id,
-            session_id="",  # session_id is not used in the adapter delete method
+        logger.info(
+            f"Return {len(memory_events)} memory events for query: {query} index={index}"
         )
+        return SearchMemoryResponse(memories=memory_events)

veadk/runner.py

@@ -70,6 +70,10 @@ class Runner:
             memory_service=self.long_term_memory,
         )
 
+        if getattr(self.agent, "tracers", None):
+            for tracers in self.agent.tracers:
+                tracers.set_app_name(self.app_name)
+
     def _convert_messages(self, messages) -> list:
         if isinstance(messages, str):
             messages = [types.Content(role="user", parts=[types.Part(text=messages)])]

veadk/tools/builtin_tools/las.py

@@ -12,12 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset, SseConnectionParams
+from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
 
 from veadk.config import getenv
+from veadk.utils.mcp_utils import get_mcp_params
 
 url = getenv("TOOL_LAS_URL")
 dataset_id = getenv("TOOL_LAS_DATASET_ID")
 
 
-las = MCPToolset(connection_params=SseConnectionParams(url=url))
+las = MCPToolset(connection_params=get_mcp_params(url=url))

veadk/{cloud/template/app.py → tools/builtin_tools/mcp_router.py}

@@ -12,19 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
+from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
 
-from config import AGENT, APP_NAME, SHORT_TERM_MEMORY, TRACERS
+from veadk.config import getenv
+from veadk.utils.mcp_utils import get_mcp_params
 
-from veadk.a2a.ve_a2a_server import init_app
+url = getenv("TOOL_MCP_ROUTER_URL")
 
-SERVER_HOST = os.getenv("SERVER_HOST")
-
-AGENT.tracers = TRACERS
-
-app = init_app(
-    server_url=SERVER_HOST,
-    app_name=APP_NAME,
-    agent=AGENT,
-    short_term_memory=SHORT_TERM_MEMORY,
-)
+mcp_router = MCPToolset(connection_params=get_mcp_params(url=url))

veadk/tools/load_knowledgebase_tool.py

@@ -51,11 +51,11 @@ class SearchKnowledgebaseResponse(BaseModel):
 
 
 async def search_knowledgebase(
-    self, query: str, app_name: str, user_id: str
+    self, query: str, app_name: str
 ) -> SearchKnowledgebaseResponse:
     """Searches the knowledgebase of the current user."""
     if isinstance(knowledgebase, KnowledgeBase):
-        res = knowledgebase.search(query, app_name=app_name, user_id=user_id)
+        res = knowledgebase.search(query, app_name=app_name)
         entry_list = []
         for r in res:
             entry_list.append(KnowledgebaseEntry(content=r))
@@ -86,9 +86,7 @@ async def load_knowledgebase(
       A list of knowledgebase results.
     """
     search_knowledgebase_response = await tool_context.search_knowledgebase(
-        query,
-        tool_context._invocation_context.app_name,
-        tool_context._invocation_context.user_id,
+        query, tool_context._invocation_context.app_name
     )
     return LoadKnowledgebaseResponse(
         knowledges=search_knowledgebase_response.knowledges

veadk/tracing/base_tracer.py

@@ -14,10 +14,11 @@
 
 import json
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Optional, Any
 
 from google.adk.agents.callback_context import CallbackContext
 from google.adk.models import LlmRequest, LlmResponse
+from google.adk.tools import BaseTool, ToolContext
 from opentelemetry import trace
 
 from veadk.utils.logger import get_logger
@@ -27,20 +28,24 @@ logger = get_logger(__name__)
 
 class BaseTracer(ABC):
     def __init__(self, name: str):
+        self.app_name = "veadk_app_name"
         pass
 
     @abstractmethod
     def dump(self) -> str: ...
 
-    def llm_metrics_hook(
+    def tracer_hook_before_model(
         self, callback_context: CallbackContext, llm_request: LlmRequest
     ) -> Optional[LlmResponse]:
+        """agent run stage"""
         trace.get_tracer("gcp.vertex.agent")
         span = trace.get_current_span()
         # logger.debug(f"llm_request: {llm_request}")
 
         req = llm_request.model_dump()
 
+        app_name = getattr(self, "app_name", "veadk_app")
+        agent_name = callback_context.agent_name
         model_name = req.get("model", "unknown")
         max_tokens = (
             None
@@ -59,6 +64,8 @@ class BaseTracer(ABC):
         )
 
         attributes = {}
+        attributes["agent.name"] = agent_name
+        attributes["app.name"] = app_name
         attributes["gen_ai.system"] = "veadk"
         if model_name:
             attributes["gen_ai.request.model"] = model_name
@@ -84,9 +91,10 @@ class BaseTracer(ABC):
             for k, v in attributes.items():
                 span.set_attribute(k, v)
 
-    def token_metrics_hook(
+    def tracer_hook_after_model(
         self, callback_context: CallbackContext, llm_response: LlmResponse
     ) -> Optional[LlmResponse]:
+        """call llm stage"""
         trace.get_tracer("gcp.vertex.agent")
         span = trace.get_current_span()
         # logger.debug(f"llm_response: {llm_response}")
@@ -95,6 +103,10 @@ class BaseTracer(ABC):
         # Refined: collect all attributes, use set_attributes, print for debugging
         attributes = {}
 
+        app_name = getattr(self, "app_name", "veadk_app")
+        agent_name = callback_context.agent_name
+        attributes["agent.name"] = agent_name
+        attributes["app.name"] = app_name
         # prompt
         user_content = callback_context.user_content
         if getattr(user_content, "role", None):
@@ -170,3 +182,32 @@ class BaseTracer(ABC):
             # Fallback for OpenTelemetry versions without set_attributes
             for k, v in attributes.items():
                 span.set_attribute(k, v)
+
+    def tracer_hook_after_tool(
+        self,
+        tool: BaseTool,
+        args: dict[str, Any],
+        tool_context: ToolContext,
+        tool_response: dict,
+    ):
+        trace.get_tracer("gcp.vertex.agent")
+        span = trace.get_current_span()
+        agent_name = tool_context.agent_name
+        tool_name = tool.name
+        app_name = getattr(self, "app_name", "veadk_app")
+        attributes = {
+            "agent.name": agent_name,
+            "app.name": app_name,
+            "tool.name": tool_name,
+        }
+
+        # Set all attributes at once if possible, else fallback to individual
+        if hasattr(span, "set_attributes"):
+            span.set_attributes(attributes)
+        else:
+            # Fallback for OpenTelemetry versions without set_attributes
+            for k, v in attributes.items():
+                span.set_attribute(k, v)
+
+    def set_app_name(self, app_name):
+        self.app_name = app_name

veadk/tracing/telemetry/exporters/inmemory_exporter.py

@@ -83,6 +83,3 @@ class InMemoryExporter(BaseModel, BaseExporter):
         self._real_exporter = self._exporter
         processor = export.SimpleSpanProcessor(self._exporter)
         return processor, None
-
-    def get_meter_context(self) -> Any:
-        return None