170 lines
5.1 KiB
Python
170 lines
5.1 KiB
Python
"""
|
||
Dev Agent 命令行交互工具
|
||
使用配置好的 Runner（SQLite 会话持久化 + Memory + 上下文压缩），
|
||
退出后再次进入同一个 session_id 可以继续对话。
|
||
|
||
使用方式：
|
||
python chat.py # 新会话，自动生成 session_id
|
||
python chat.py --session my_session # 指定 session_id
|
||
python chat.py --list # 列出所有会话
|
||
python chat.py --delete my_session # 删除某个会话
|
||
"""
|
||
import os
|
||
import sys
|
||
import asyncio
|
||
import argparse
|
||
|
||
# 确保项目根目录在 path 里
|
||
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
|
||
if PROJECT_ROOT not in sys.path:
|
||
sys.path.insert(0, PROJECT_ROOT)
|
||
|
||
from dotenv import load_dotenv
|
||
load_dotenv(os.path.join(PROJECT_ROOT, "my_agent", ".env"))
|
||
|
||
# 强制 UTF-8
|
||
os.environ["PYTHONUTF8"] = "1"
|
||
|
||
from google.adk.runners import Runner
|
||
from google.adk.sessions.sqlite_session_service import SqliteSessionService
|
||
from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
|
||
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
|
||
from google.genai import types
|
||
from my_agent.app import dev_app
|
||
|
||
|
||
# 数据目录
|
||
DATA_DIR = os.path.join(PROJECT_ROOT, "data")
|
||
os.makedirs(DATA_DIR, exist_ok=True)
|
||
|
||
# 同一个数据库，A2A server 和 CLI 共享
|
||
DB_PATH = os.path.join(DATA_DIR, "sessions.db")
|
||
|
||
USER_ID = "local_user"
|
||
|
||
|
||
def get_runner() -> Runner:
|
||
"""创建带 SQLite 会话持久化的 Runner"""
|
||
session_service = SqliteSessionService(db_path=DB_PATH)
|
||
memory_service = InMemoryMemoryService()
|
||
artifact_service = InMemoryArtifactService()
|
||
|
||
return Runner(
|
||
app=dev_app,
|
||
session_service=session_service,
|
||
memory_service=memory_service,
|
||
artifact_service=artifact_service,
|
||
auto_create_session=True,
|
||
)
|
||
|
||
|
||
async def list_sessions():
|
||
"""列出所有会话"""
|
||
runner = get_runner()
|
||
response = await runner.session_service.list_sessions(
|
||
app_name=dev_app.name,
|
||
user_id=USER_ID,
|
||
)
|
||
sessions = response.sessions
|
||
if not sessions:
|
||
print("（暂无会话）")
|
||
return
|
||
|
||
print(f"共 {len(sessions)} 个会话：\n")
|
||
for s in sessions:
|
||
# 取第一条用户消息作为摘要
|
||
preview = ""
|
||
for e in s.events:
|
||
if e.content and e.content.parts and e.author == "user":
|
||
text = e.content.parts[0].text[:50]
|
||
preview = f" → {text}"
|
||
break
|
||
print(f" [{s.id}] {preview}")
|
||
print(f" 更新时间: {s.last_update_time}")
|
||
|
||
|
||
async def delete_session(session_id: str):
|
||
"""删除指定会话"""
|
||
runner = get_runner()
|
||
try:
|
||
await runner.session_service.delete_session(
|
||
app_name=dev_app.name,
|
||
user_id=USER_ID,
|
||
session_id=session_id,
|
||
)
|
||
print(f"会话 [{session_id}] 已删除")
|
||
except Exception as e:
|
||
print(f"删除失败: {e}")
|
||
|
||
|
||
async def chat(session_id: str | None = None):
|
||
"""交互式对话"""
|
||
runner = get_runner()
|
||
|
||
# 如果没有指定 session_id，自动创建
|
||
if not session_id:
|
||
session = await runner.session_service.create_session(
|
||
app_name=dev_app.name,
|
||
user_id=USER_ID,
|
||
)
|
||
session_id = session.id
|
||
print(f"新会话已创建，session_id: {session_id}")
|
||
print(f"下次可用 `python chat.py --session {session_id}` 继续\n")
|
||
|
||
print(f"=== Dev Agent 对话 ===")
|
||
print(f"Session: {session_id}")
|
||
print(f"输入消息开始对话，输入 quit / exit 退出\n")
|
||
|
||
while True:
|
||
try:
|
||
user_input = input("你: ").strip()
|
||
except (EOFError, KeyboardInterrupt):
|
||
print("\n再见！")
|
||
break
|
||
|
||
if not user_input:
|
||
continue
|
||
if user_input.lower() in ("quit", "exit", "退出"):
|
||
print("再见！")
|
||
break
|
||
|
||
print("花花: ", end="", flush=True)
|
||
|
||
try:
|
||
full_response = ""
|
||
async for event in runner.run_async(
|
||
user_id=USER_ID,
|
||
session_id=session_id,
|
||
new_message=types.Content(parts=[types.Part(text=user_input)]),
|
||
):
|
||
if event.is_final_response():
|
||
# 最终回复
|
||
for part in event.content.parts:
|
||
if hasattr(part, "text") and part.text:
|
||
print(part.text, end="", flush=True)
|
||
full_response += part.text
|
||
print()
|
||
except Exception as e:
|
||
print(f"\n[出错] {e}")
|
||
|
||
print()
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Dev Agent 命令行交互工具")
|
||
parser.add_argument("--session", "-s", help="会话 ID，指定后继续该会话")
|
||
parser.add_argument("--list", "-l", action="store_true", help="列出所有会话")
|
||
parser.add_argument("--delete", "-d", help="删除指定会话")
|
||
args = parser.parse_args()
|
||
|
||
if args.list:
|
||
asyncio.run(list_sessions())
|
||
elif args.delete:
|
||
asyncio.run(delete_session(args.delete))
|
||
else:
|
||
asyncio.run(chat(args.session))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|