299 lines
11 KiB
Python
299 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Vertex AI Memory Bank Web界面
|
||
一个简单的Streamlit应用，用于通过Web界面访问和管理Memory Bank
|
||
"""
|
||
|
||
import streamlit as st
|
||
import asyncio
|
||
import sys
|
||
import os
|
||
from datetime import datetime
|
||
|
||
# 添加项目路径
|
||
sys.path.append('/Users/ben/liurenchaxin/src')
|
||
|
||
try:
|
||
from jixia.memory.factory import get_memory_backend
|
||
except ImportError as e:
|
||
st.error(f"无法导入jixia模块: {e}")
|
||
st.info("请确保已激活虚拟环境并安装了所需依赖")
|
||
st.stop()
|
||
|
||
# 页面配置
|
||
st.set_page_config(
|
||
page_title="Memory Bank 管理界面",
|
||
page_icon="🧠",
|
||
layout="wide"
|
||
)
|
||
|
||
# 标题
|
||
st.title("🧠 Vertex AI Memory Bank 管理界面")
|
||
st.markdown("---")
|
||
|
||
# 侧边栏配置
|
||
st.sidebar.header("配置")
|
||
project_id = st.sidebar.text_input("项目ID", value="inner-radius-469712-e9")
|
||
location = st.sidebar.text_input("区域", value="us-central1")
|
||
|
||
# 八仙列表
|
||
EIGHT_IMMORTALS = [
|
||
"lvdongbin", "tieguaili", "hanxiangzi", "lanzaihe",
|
||
"hesengu", "zhonghanli", "caogujiu", "hanzhongli"
|
||
]
|
||
|
||
# 缓存Memory Bank后端
|
||
@st.cache_resource
|
||
def get_memory_backend_cached():
|
||
"""获取Memory Bank后端（缓存）"""
|
||
try:
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
backend = loop.run_until_complete(get_memory_backend("vertex"))
|
||
return backend, loop
|
||
except Exception as e:
|
||
st.error(f"初始化Memory Bank失败: {e}")
|
||
return None, None
|
||
|
||
# 异步函数包装器
|
||
def run_async(coro):
|
||
"""运行异步函数"""
|
||
backend, loop = get_memory_backend_cached()
|
||
if backend is None:
|
||
return None
|
||
try:
|
||
return loop.run_until_complete(coro)
|
||
except Exception as e:
|
||
st.error(f"操作失败: {e}")
|
||
return None
|
||
|
||
# 主界面
|
||
tab1, tab2, tab3, tab4 = st.tabs(["📋 Memory Bank列表", "🔍 搜索记忆", "➕ 添加记忆", "📊 统计信息"])
|
||
|
||
with tab1:
|
||
st.header("Memory Bank 实例列表")
|
||
|
||
if st.button("🔄 刷新列表", key="refresh_list"):
|
||
st.rerun()
|
||
|
||
# 显示八仙Memory Bank状态
|
||
cols = st.columns(4)
|
||
for i, immortal in enumerate(EIGHT_IMMORTALS):
|
||
with cols[i % 4]:
|
||
with st.container():
|
||
st.subheader(f"🧙♂️ {immortal}")
|
||
|
||
# 检查Memory Bank状态
|
||
backend, _ = get_memory_backend_cached()
|
||
if backend:
|
||
try:
|
||
# 尝试获取agent context来验证Memory Bank存在
|
||
context = run_async(backend.get_agent_context(immortal))
|
||
if context is not None:
|
||
st.success("✅ 活跃")
|
||
|
||
# 显示记忆数量
|
||
memories = run_async(backend.search_memories(immortal, "", limit=100))
|
||
if memories:
|
||
st.info(f"📝 记忆数量: {len(memories)}")
|
||
else:
|
||
st.info("📝 记忆数量: 0")
|
||
else:
|
||
st.warning("⚠️ 未初始化")
|
||
except Exception as e:
|
||
st.error(f"❌ 错误: {str(e)[:50]}...")
|
||
else:
|
||
st.error("❌ 连接失败")
|
||
|
||
with tab2:
|
||
st.header("🔍 搜索记忆")
|
||
|
||
col1, col2 = st.columns([1, 2])
|
||
|
||
with col1:
|
||
selected_agent = st.selectbox("选择Agent", EIGHT_IMMORTALS, key="search_agent")
|
||
search_query = st.text_input("搜索关键词", placeholder="输入要搜索的内容...", key="search_query")
|
||
search_limit = st.slider("结果数量", 1, 50, 10, key="search_limit")
|
||
|
||
if st.button("🔍 搜索", key="search_button"):
|
||
if search_query:
|
||
with st.spinner("搜索中..."):
|
||
backend, _ = get_memory_backend_cached()
|
||
if backend:
|
||
memories = run_async(backend.search_memories(selected_agent, search_query, limit=search_limit))
|
||
st.session_state['search_results'] = memories
|
||
st.session_state['search_agent'] = selected_agent
|
||
st.session_state['search_query'] = search_query
|
||
else:
|
||
st.warning("请输入搜索关键词")
|
||
|
||
with col2:
|
||
st.subheader("搜索结果")
|
||
|
||
if 'search_results' in st.session_state and st.session_state['search_results']:
|
||
st.success(f"找到 {len(st.session_state['search_results'])} 条记忆")
|
||
|
||
for i, memory in enumerate(st.session_state['search_results']):
|
||
with st.expander(f"记忆 {i+1}: {memory.get('content', 'N/A')[:50]}..."):
|
||
st.write(f"**内容**: {memory.get('content', 'N/A')}")
|
||
st.write(f"**类型**: {memory.get('memory_type', 'N/A')}")
|
||
st.write(f"**时间**: {memory.get('timestamp', 'N/A')}")
|
||
if 'metadata' in memory:
|
||
st.write(f"**元数据**: {memory['metadata']}")
|
||
elif 'search_results' in st.session_state:
|
||
st.info("未找到匹配的记忆")
|
||
else:
|
||
st.info("请执行搜索以查看结果")
|
||
|
||
with tab3:
|
||
st.header("➕ 添加记忆")
|
||
|
||
col1, col2 = st.columns([1, 1])
|
||
|
||
with col1:
|
||
add_agent = st.selectbox("选择Agent", EIGHT_IMMORTALS, key="add_agent")
|
||
memory_type = st.selectbox("记忆类型", ["conversation", "preference", "knowledge", "strategy"], key="memory_type")
|
||
memory_content = st.text_area("记忆内容", placeholder="输入要添加的记忆内容...", height=150, key="memory_content")
|
||
|
||
# 可选的元数据
|
||
st.subheader("元数据（可选）")
|
||
importance = st.slider("重要性", 1, 10, 5, key="importance")
|
||
tags = st.text_input("标签（用逗号分隔）", placeholder="标签1, 标签2, 标签3", key="tags")
|
||
|
||
if st.button("➕ 添加记忆", key="add_memory_button"):
|
||
if memory_content:
|
||
with st.spinner("添加记忆中..."):
|
||
backend, _ = get_memory_backend_cached()
|
||
if backend:
|
||
# 准备元数据
|
||
metadata = {
|
||
"importance": importance,
|
||
"timestamp": datetime.now().isoformat(),
|
||
"source": "web_interface"
|
||
}
|
||
if tags:
|
||
metadata["tags"] = [tag.strip() for tag in tags.split(",")]
|
||
|
||
# 添加记忆
|
||
success = run_async(backend.add_memory(
|
||
agent_id=add_agent,
|
||
content=memory_content,
|
||
memory_type=memory_type,
|
||
metadata=metadata
|
||
))
|
||
|
||
if success:
|
||
st.success("✅ 记忆添加成功！")
|
||
# 清空输入
|
||
st.session_state['memory_content'] = ""
|
||
st.session_state['tags'] = ""
|
||
else:
|
||
st.error("❌ 添加记忆失败")
|
||
else:
|
||
st.warning("请输入记忆内容")
|
||
|
||
with col2:
|
||
st.subheader("添加记忆预览")
|
||
if memory_content:
|
||
st.info(f"**Agent**: {add_agent}")
|
||
st.info(f"**类型**: {memory_type}")
|
||
st.info(f"**内容**: {memory_content}")
|
||
st.info(f"**重要性**: {importance}/10")
|
||
if tags:
|
||
st.info(f"**标签**: {tags}")
|
||
else:
|
||
st.info("输入记忆内容以查看预览")
|
||
|
||
with tab4:
|
||
st.header("📊 统计信息")
|
||
|
||
if st.button("🔄 刷新统计", key="refresh_stats"):
|
||
st.rerun()
|
||
|
||
# 获取统计信息
|
||
backend, _ = get_memory_backend_cached()
|
||
if backend:
|
||
stats_data = []
|
||
|
||
for immortal in EIGHT_IMMORTALS:
|
||
try:
|
||
# 获取记忆数量
|
||
memories = run_async(backend.search_memories(immortal, "", limit=1000))
|
||
memory_count = len(memories) if memories else 0
|
||
|
||
# 获取agent context
|
||
context = run_async(backend.get_agent_context(immortal))
|
||
status = "活跃" if context else "未初始化"
|
||
|
||
stats_data.append({
|
||
"Agent": immortal,
|
||
"状态": status,
|
||
"记忆数量": memory_count,
|
||
"最后更新": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
})
|
||
except Exception as e:
|
||
stats_data.append({
|
||
"Agent": immortal,
|
||
"状态": "错误",
|
||
"记忆数量": 0,
|
||
"最后更新": f"错误: {str(e)[:30]}..."
|
||
})
|
||
|
||
# 显示统计表格
|
||
st.dataframe(stats_data, use_container_width=True)
|
||
|
||
# 显示汇总信息
|
||
col1, col2, col3, col4 = st.columns(4)
|
||
|
||
total_agents = len(EIGHT_IMMORTALS)
|
||
active_agents = sum(1 for item in stats_data if item["状态"] == "活跃")
|
||
total_memories = sum(item["记忆数量"] for item in stats_data)
|
||
avg_memories = total_memories / total_agents if total_agents > 0 else 0
|
||
|
||
with col1:
|
||
st.metric("总Agent数", total_agents)
|
||
|
||
with col2:
|
||
st.metric("活跃Agent数", active_agents)
|
||
|
||
with col3:
|
||
st.metric("总记忆数", total_memories)
|
||
|
||
with col4:
|
||
st.metric("平均记忆数", f"{avg_memories:.1f}")
|
||
|
||
# 页脚
|
||
st.markdown("---")
|
||
st.markdown(
|
||
"""
|
||
<div style='text-align: center; color: #666;'>
|
||
🧠 Vertex AI Memory Bank Web界面 |
|
||
<a href='https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/memory-bank/overview' target='_blank'>官方文档</a>
|
||
</div>
|
||
""",
|
||
unsafe_allow_html=True
|
||
)
|
||
|
||
# 使用说明
|
||
with st.expander("📖 使用说明"):
|
||
st.markdown("""
|
||
### 功能说明
|
||
|
||
1. **Memory Bank列表**: 查看所有八仙角色的Memory Bank状态和记忆数量
|
||
2. **搜索记忆**: 在指定Agent的记忆中搜索特定内容
|
||
3. **添加记忆**: 为Agent添加新的记忆，支持不同类型和元数据
|
||
4. **统计信息**: 查看所有Agent的统计数据和汇总信息
|
||
|
||
### 使用前准备
|
||
|
||
1. 确保已激活虚拟环境: `source venv/bin/activate`
|
||
2. 确保已设置Google Cloud认证: `gcloud auth application-default login`
|
||
3. 运行此界面: `streamlit run memory_bank_web_interface.py`
|
||
|
||
### 注意事项
|
||
|
||
- Memory Bank目前仅在us-central1区域可用
|
||
- 搜索功能支持模糊匹配
|
||
- 添加的记忆会立即生效
|
||
- 统计信息实时更新
|
||
""") |