from veadk import Agent, Runner
from veadk.knowledgebase import KnowledgeBase
from veadk.memory import LongTermMemory, ShortTermMemory
from veadk.tools.builtin_tools.run_code import run_code
from veadk.tools.builtin_tools.mcp_router import mcp_router
from google.adk.agents.callback_context import CallbackContext
from dotenv import load_dotenv
memory_name = os.getenv("DATABASE_VIKINGMEM_COLLECTION")
raise ValueError("DATABASE_VIKINGMEM_COLLECTION environment variable is not set")
long_term_memory = LongTermMemory(backend="viking", index=memory_name)
short_term_memory = ShortTermMemory(backend="local")
knowledge_name = os.getenv("DATABASE_VIKING_COLLECTION")
raise ValueError("DATABASE_VIKING_COLLECTION environment variable is not set")
knowledgebase = KnowledgeBase(backend="viking", index=knowledge_name)
async def after_agent_execution(callback_context: CallbackContext):
session = callback_context._invocation_context.session
await long_term_memory.add_session_to_memory(session)
personal_assistant = Agent(
name="personal_assistant",
description="一个智能个人助手,可以帮助用户管理日常事务和提供信息。",
instruction="""你是一个友好的智能个人助手,请用简洁明了的语言回答用户的问题。
如果用户询问与知识库相关的信息,请使用知识库检索。
如果用户提到个人偏好或重要信息,请将其存储到长期记忆中。
short_term_memory=short_term_memory,
long_term_memory=long_term_memory,
knowledgebase=knowledgebase,
tools=[run_code, mcp_router],
after_agent_callback=after_agent_execution,
root_agent = personal_assistant
print("输入 'exit' 或 'quit' 退出程序。")
agent=personal_assistant,
app_name="personal_assistant_demo",
short_term_memory=short_term_memory
user_input = input("请输入您的问题: ")
if user_input.lower() in ["exit", "quit", "退出"]:
response = await runner.run(messages=user_input)
print(f"助手: {response}")
if __name__ == "__main__":