diff --git a/veadk/knowledgebase/knowledgebase.py b/veadk/knowledgebase/knowledgebase.py index 98890aad..854c8d66 100644 --- a/veadk/knowledgebase/knowledgebase.py +++ b/veadk/knowledgebase/knowledgebase.py @@ -108,6 +108,8 @@ class KnowledgeBase(BaseModel): enable_profile: bool = False + query_with_user_profile: bool = False + def model_post_init(self, __context: Any) -> None: if isinstance(self.backend, BaseKnowledgebaseBackend): self._backend = self.backend @@ -134,6 +136,11 @@ def model_post_init(self, __context: Any) -> None: f"Initialized knowledgebase with backend {self._backend.__class__.__name__}" ) + if self.query_with_user_profile: + logger.info( + "Enable user profile querying for knowledgebase. You *must* use Viking Memory backend to enjoy this feature." + ) + def add_from_directory(self, directory: str, **kwargs) -> bool: """Add knowledge from file path to knowledgebase. diff --git a/veadk/memory/long_term_memory.py b/veadk/memory/long_term_memory.py index cd0c8486..7f64a8c9 100644 --- a/veadk/memory/long_term_memory.py +++ b/veadk/memory/long_term_memory.py @@ -100,68 +100,6 @@ class LongTermMemory(BaseMemoryService, BaseModel): Notes: Please ensure that you have set the embedding-related configurations in environment variables. - - Examples: - ### Simple long-term memory - - Once create a long-term memory withou any arguments, all configurations are come from **environment variables**. - - ```python - import asyncio - - from veadk import Agent, Runner - from veadk.memory.long_term_memory import LongTermMemory - from veadk.memory.short_term_memory import ShortTermMemory - - app_name = "veadk_playground_app" - user_id = "veadk_playground_user" - - long_term_memory = LongTermMemory(backend="local", app_name=app_name) - - agent = Agent(long_term_memory=long_term_memory) - - runner = Runner( - agent=agent, - app_name=app_name, - user_id=user_id, - short_term_memory=ShortTermMemory(), - ) - - # ===== add memory ===== - session_id = "veadk_playground_session" - teaching_prompt = "I brought an ice-cream last week." - - asyncio.run(runner.run(messages=teaching_prompt, session_id=session_id)) - asyncio.run( - runner.save_session_to_long_term_memory(session_id=session_id) - ) # save session to long-term memory - - - # ===== check memory ===== - session_id = "veadk_playground_session_2" # use a new session - student_prompt = "What I brought last week?" - - response = asyncio.run(runner.run(messages=student_prompt, session_id=session_id)) - - print(response) - ``` - - ### Create with a backend instance - - ```python - from veadk.memory.long_term_memory import LongTermMemory - from veadk.memory.long_term_memory.backends import LongTermMemory - - long_term_memory = LongTermMemory(backend=...) - ``` - - ### Create with backend configurations - - ```python - from veadk.memory.long_term_memory import LongTermMemory - - long_term_memory = LongTermMemory(backend="", backend_config={}) - ``` """ backend: Union[ @@ -313,20 +251,6 @@ async def search_memory( SearchMemoryResponse: An object containing a list of `MemoryEntry` items representing the retrieved memory snippets relevant to the query. - - Examples: - ```python - response = await memory_service.search_memory( - app_name="chat_app", - user_id="user_123", - query="favorite programming language" - ) - - for memory in response.memories: - print(memory.content.parts[0].text) - # Output: - # User likes Python and TypeScript for backend development. - ``` """ logger.info(f"Search memory with query={query}") @@ -369,3 +293,13 @@ async def search_memory( f"Return {len(memory_events)} memory events for query: {query} index={self.index} user_id={user_id}" ) return SearchMemoryResponse(memories=memory_events) + + def get_user_profile(self, user_id: str) -> str: + logger.info(f"Get user profile for user_id={user_id}") + if self.backend == "viking": + return self._backend.get_user_profile(user_id=user_id) # type: ignore + else: + logger.error( + f"Long term memory backend {self.backend} does not support get user profile. Return empty string." + ) + return "" diff --git a/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py b/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py index 361b1772..9c9d415d 100644 --- a/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py +++ b/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py @@ -224,3 +224,58 @@ def search_memory( for r in result ] return [] + + def get_user_profile(self, user_id: str) -> str: + from veadk.utils.volcengine_sign import ve_request + + response: dict = self._get_client().get_collection( + collection_name=self.index, project=self.volcengine_project + ) + + mem_id = response["Result"]["ResourceId"] + logger.info( + f"Get user profile for user_id={user_id} from Viking Memory with mem_id={mem_id}" + ) + + ak = "" + sk = "" + sts_token = "" + if self.volcengine_access_key and self.volcengine_secret_key: + ak = self.volcengine_access_key + sk = self.volcengine_secret_key + sts_token = self.session_token + else: + cred = get_credential_from_vefaas_iam() + ak = cred.access_key_id + sk = cred.secret_access_key + sts_token = cred.session_token + + response = ve_request( + request_body={ + "filter": { + "user_id": [user_id], + "memory_category": 1, + }, + "limit": 5000, + "resource_id": mem_id, + }, + action="MemorySearch", + ak=ak, + sk=sk, + header={"X-Security-Token": sts_token}, + service="vikingdb", + version="2025-06-09", + region=self.region, + host="open.volcengineapi.com", + ) + + try: + logger.debug( + f"Response from VikingDB: {response}, user_profile: {response['data']['result_list'][0]['memory_info']['user_profile']}" + ) + return response["data"]["result_list"][0]["memory_info"]["user_profile"] + except (KeyError, IndexError): + logger.error( + f"Failed to get user profile for user_id={user_id} mem_id={mem_id}: {response}" + ) + return "" diff --git a/veadk/tools/builtin_tools/load_knowledgebase.py b/veadk/tools/builtin_tools/load_knowledgebase.py index 52cad6ec..51b80944 100644 --- a/veadk/tools/builtin_tools/load_knowledgebase.py +++ b/veadk/tools/builtin_tools/load_knowledgebase.py @@ -124,6 +124,31 @@ async def process_llm_request( ] ) + 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: