278 lines
9.6 KiB
Python
278 lines
9.6 KiB
Python
"""
|
||
Luna 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
|
||
|
||
# 脚本所在目录（作为 .env / data 等相对路径的基准）
|
||
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
|
||
if PROJECT_ROOT not in sys.path:
|
||
sys.path.insert(0, PROJECT_ROOT)
|
||
|
||
# 项目根目录（往上两级），确保 from agents.xxx.xxx import 可用
|
||
_REPO_ROOT = os.path.abspath(os.path.join(PROJECT_ROOT, "../.."))
|
||
if _REPO_ROOT not in sys.path:
|
||
sys.path.insert(0, _REPO_ROOT)
|
||
|
||
from dotenv import load_dotenv
|
||
load_dotenv(os.path.join(PROJECT_ROOT, "", ".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.artifacts.in_memory_artifact_service import InMemoryArtifactService
|
||
from google.adk.agents.run_config import RunConfig, StreamingMode
|
||
from google.genai import types
|
||
from agents.luna.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_luna.db")
|
||
|
||
USER_ID = "local_user"
|
||
|
||
|
||
def get_runner() -> Runner:
|
||
"""创建带 SQLite 会话持久化的 Runner"""
|
||
session_service = SqliteSessionService(db_path=DB_PATH)
|
||
artifact_service = InMemoryArtifactService()
|
||
|
||
return Runner(
|
||
app=dev_app,
|
||
session_service=session_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"=== Luna Agent 对话 ===")
|
||
print(f"模型: {dev_app.root_agent.model.model}")
|
||
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("Luna: ", end="", flush=True)
|
||
|
||
async def _agent_task():
|
||
"""运行 agent 并流式输出，返回是否完成"""
|
||
displayed_text = ""
|
||
thought_printed = False
|
||
run_config = RunConfig(streaming_mode=StreamingMode.SSE)
|
||
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)]),
|
||
run_config=run_config,
|
||
):
|
||
if not event.content or not event.content.parts:
|
||
continue
|
||
|
||
parts = event.content.parts
|
||
|
||
# 1. 思考内容（thought parts）——灰色流式显示
|
||
thought_parts = [
|
||
p.text for p in parts
|
||
if hasattr(p, "text") and p.text
|
||
and getattr(p, "thought", False)
|
||
]
|
||
if thought_parts:
|
||
thought_text = "".join(thought_parts)
|
||
if not thought_printed:
|
||
print("\n\033[90m（思考中…", end="", flush=True)
|
||
thought_printed_outer[0] = True
|
||
thought_displayed_outer[0] = 0
|
||
if len(thought_text) > thought_displayed_outer[0]:
|
||
print(thought_text[thought_displayed_outer[0]:], end="", flush=True)
|
||
thought_displayed_outer[0] = len(thought_text)
|
||
|
||
# 2. 正式文本——增量显示
|
||
text_parts = [
|
||
p.text for p in parts
|
||
if hasattr(p, "text") and p.text
|
||
and not getattr(p, "thought", False)
|
||
]
|
||
if text_parts:
|
||
text = "".join(text_parts)
|
||
if len(text) > len(displayed_text):
|
||
if thought_printed_outer[0]:
|
||
print("\033[0m\nLuna: ", end="", flush=True)
|
||
thought_printed_outer[0] = False
|
||
new_text = text[len(displayed_text):]
|
||
print(new_text, end="", flush=True)
|
||
displayed_text = text
|
||
|
||
# 3. 工具调用提示
|
||
fcalls = event.get_function_calls()
|
||
if fcalls and not event.partial:
|
||
if thought_printed_outer[0]:
|
||
print("\033[0m", end="", flush=True)
|
||
thought_printed_outer[0] = False
|
||
for fc in fcalls:
|
||
args_str = str(fc.args)[:80]
|
||
print(f"\n\033[36m🔧 调用工具: {fc.name}({args_str})\033[0m")
|
||
print("Luna: ", end="", flush=True)
|
||
|
||
# 4. 最终响应
|
||
if event.is_final_response() and not event.partial:
|
||
if thought_printed_outer[0]:
|
||
print("\033[0m", end="", flush=True)
|
||
thought_printed_outer[0] = False
|
||
print()
|
||
return True
|
||
return False
|
||
|
||
thought_printed_outer = [False]
|
||
thought_displayed_outer = [0]
|
||
|
||
# 启动 agent 任务 + 按键监听
|
||
task = asyncio.create_task(_agent_task())
|
||
|
||
async def _keyboard_listener():
|
||
"""监听按键，检测到中断键时取消 agent 任务"""
|
||
if sys.platform != "win32":
|
||
return
|
||
import msvcrt
|
||
while not task.done():
|
||
await asyncio.sleep(0.05)
|
||
if msvcrt.kbhit():
|
||
ch = msvcrt.getwch()
|
||
# 支持的中断键: Ctrl+C (0x03), Esc (0x1b), q/Q
|
||
if ch in ("\x03", "\x1b", "q", "Q"):
|
||
task.cancel()
|
||
return
|
||
# 功能键/方向键是两个字节的，跳过第二个
|
||
if ch in ("\xe0", "\x00"):
|
||
try:
|
||
msvcrt.getwch()
|
||
except Exception:
|
||
pass
|
||
|
||
try:
|
||
kb_task = asyncio.create_task(_keyboard_listener())
|
||
await task
|
||
kb_task.cancel()
|
||
try:
|
||
await kb_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
except asyncio.CancelledError:
|
||
if thought_printed_outer[0]:
|
||
print("\033[0m", end="")
|
||
print("\n\033[33m[已中断] 按回车继续输入新指令\033[0m")
|
||
# 清空可能残留的输入缓冲
|
||
if sys.platform == "win32":
|
||
import msvcrt
|
||
while msvcrt.kbhit():
|
||
msvcrt.getwch()
|
||
try:
|
||
input()
|
||
except EOFError:
|
||
pass
|
||
continue
|
||
except Exception as e:
|
||
if thought_printed_outer[0]:
|
||
print("\033[0m", end="")
|
||
print(f"\n[出错] {e}")
|
||
|
||
print()
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Luna 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()
|