Improve scripts

chat.py

@@ -1,19 +1,27 @@
 #!/usr/bin/env python3
 """Minimal CLI chat agent to test FirestoreSessionService."""
 
+import argparse
 import asyncio
+import logging
 
 from google import genai
 from google.adk.agents import LlmAgent
 from google.adk.runners import Runner
 from google.cloud.firestore_v1.async_client import AsyncClient
 from google.genai import types
+from rich.console import Console
+from rich.logging import RichHandler
+from rich.markdown import Markdown
+from rich.panel import Panel
 
 from adk_firestore_sessionmanager import FirestoreSessionService
 
 APP_NAME = "test_agent"
 USER_ID = "dev_user"
 
+console = Console()
+
 root_agent = LlmAgent(
     name=APP_NAME,
     model="gemini-2.5-flash",
@@ -21,11 +29,29 @@ root_agent = LlmAgent(
 )
 
 
-async def main() -> None:
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Chat with a Firestore-backed ADK agent.")
+    parser.add_argument(
+        "--log-level",
+        default="WARNING",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        help="Set the log level (default: WARNING)",
+    )
+    return parser.parse_args()
+
+
+async def main(args: argparse.Namespace) -> None:
+    logging.basicConfig(
+        level=args.log_level,
+        format="%(message)s",
+        datefmt="[%X]",
+        handlers=[RichHandler(console=console, rich_tracebacks=True)],
+    )
+
     db = AsyncClient()
     session_service = FirestoreSessionService(
         db=db,
-        compaction_token_threshold=800_000,
+        compaction_token_threshold=10_000,
         genai_client=genai.Client(),
     )
 
@@ -45,19 +71,19 @@ async def main() -> None:
             user_id=USER_ID,
             session_id=resp.sessions[0].id,
         )
-        print(f"Resuming session {session.id}.")
+        console.print(f"Resuming session [bold cyan]{session.id}[/]")
     else:
         session = await session_service.create_session(
             app_name=APP_NAME,
             user_id=USER_ID,
         )
-        print(f"Session {session.id} created.")
+        console.print(f"Session [bold cyan]{session.id}[/] created.")
 
-    print("Type 'exit' to quit.\n")
+    console.print("Type [bold]exit[/] to quit.\n")
 
     while True:
         try:
-            user_input = input("You: ").strip()
+            user_input = console.input("[bold green]You:[/] ").strip()
         except (EOFError, KeyboardInterrupt):
             break
         if not user_input or user_input.lower() == "exit":
@@ -73,11 +99,11 @@ async def main() -> None:
             if event.content and event.content.parts and not event.partial:
                 text = "".join(p.text or "" for p in event.content.parts)
                 if text:
-                    print(f"Agent: {text}")
+                    console.print(Panel(Markdown(text), title="Agent", border_style="blue"))
 
     await runner.close()
-    print("\nGoodbye!")
+    console.print("\n[dim]Goodbye![/]")
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    asyncio.run(main(parse_args()))

pyproject.toml

@@ -21,6 +21,7 @@ dev = [
     "pytest>=9.0.2",
     "pytest-asyncio>=0.24",
     "pytest-cov>=7.0.0",
+    "rich>=14.0.0",
 ]
 
 [tool.pytest.ini_options]

uv.lock

@@ -21,6 +21,7 @@ dev = [
     { name = "pytest" },
     { name = "pytest-asyncio" },
     { name = "pytest-cov" },
+    { name = "rich" },
 ]
 
 [package.metadata]
@@ -34,6 +35,7 @@ dev = [
     { name = "pytest", specifier = ">=9.0.2" },
     { name = "pytest-asyncio", specifier = ">=0.24" },
     { name = "pytest-cov", specifier = ">=7.0.0" },
+    { name = "rich", specifier = ">=14.0.0" },
 ]
 
 [[package]]
@@ -1432,6 +1434,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/87/fb/99f81ac72ae23375f22b7afdb7642aba97c00a713c217124420147681a2f/mako-1.3.10-py3-none-any.whl", hash = "sha256:baef24a52fc4fc514a0887ac600f9f1cff3d82c61d4d700a1fa84d597b88db59", size = 78509, upload-time = "2025-04-10T12:50:53.297Z" },
 ]
 
+[[package]]
+name = "markdown-it-py"
+version = "4.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "mdurl" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/5b/f5/4ec618ed16cc4f8fb3b701563655a69816155e79e24a17b651541804721d/markdown_it_py-4.0.0.tar.gz", hash = "sha256:cb0a2b4aa34f932c007117b194e945bd74e0ec24133ceb5bac59009cda1cb9f3", size = 73070, upload-time = "2025-08-11T12:57:52.854Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147", size = 87321, upload-time = "2025-08-11T12:57:51.923Z" },
+]
+
 [[package]]
 name = "markupsafe"
 version = "3.0.3"
@@ -1520,6 +1534,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/fd/d9/eaa1f80170d2b7c5ba23f3b59f766f3a0bb41155fbc32a69adfa1adaaef9/mcp-1.26.0-py3-none-any.whl", hash = "sha256:904a21c33c25aa98ddbeb47273033c435e595bbacfdb177f4bd87f6dceebe1ca", size = 233615, upload-time = "2026-01-24T19:40:30.652Z" },
 ]
 
+[[package]]
+name = "mdurl"
+version = "0.1.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" },
+]
+
 [[package]]
 name = "mmh3"
 version = "5.2.0"
@@ -2352,6 +2375,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/1e/db/4254e3eabe8020b458f1a747140d32277ec7a271daf1d235b70dc0b4e6e3/requests-2.32.5-py3-none-any.whl", hash = "sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6", size = 64738, upload-time = "2025-08-18T20:46:00.542Z" },
 ]
 
+[[package]]
+name = "rich"
+version = "14.3.3"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "markdown-it-py" },
+    { name = "pygments" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/b3/c6/f3b320c27991c46f43ee9d856302c70dc2d0fb2dba4842ff739d5f46b393/rich-14.3.3.tar.gz", hash = "sha256:b8daa0b9e4eef54dd8cf7c86c03713f53241884e814f4e2f5fb342fe520f639b", size = 230582, upload-time = "2026-02-19T17:23:12.474Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/14/25/b208c5683343959b670dc001595f2f3737e051da617f66c31f7c4fa93abc/rich-14.3.3-py3-none-any.whl", hash = "sha256:793431c1f8619afa7d3b52b2cdec859562b950ea0d4b6b505397612db8d5362d", size = 310458, upload-time = "2026-02-19T17:23:13.732Z" },
+]
+
 [[package]]
 name = "rpds-py"
 version = "0.30.0"

view_summary.py

@@ -4,12 +4,17 @@
 import asyncio
 
 from google.cloud.firestore_v1.async_client import AsyncClient
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.panel import Panel
 
 from adk_firestore_sessionmanager import FirestoreSessionService
 
 APP_NAME = "test_agent"
 USER_ID = "dev_user"
 
+console = Console()
+
 
 async def main() -> None:
     db = AsyncClient()
@@ -20,7 +25,7 @@ async def main() -> None:
     )
 
     if not resp.sessions:
-        print("No sessions found.")
+        console.print("[dim]No sessions found.[/]")
         return
 
     for s in resp.sessions:
@@ -29,12 +34,10 @@ async def main() -> None:
         data = snap.to_dict() or {}
         summary = data.get("conversation_summary")
 
-        print(f"Session: {s.id}")
         if summary:
-            print(f"Summary:\n{summary}")
+            console.print(Panel(Markdown(summary), title=f"Session {s.id}", border_style="cyan"))
         else:
-            print("No summary yet.")
-        print()
+            console.print(Panel("[dim]No summary yet.[/]", title=f"Session {s.id}", border_style="yellow"))
 
 
 if __name__ == "__main__":