diff --git a/chat.py b/chat.py index 42346a2..121f6dc 100644 --- a/chat.py +++ b/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())) diff --git a/pyproject.toml b/pyproject.toml index 37ea056..e389722 100644 --- a/pyproject.toml +++ b/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] diff --git a/uv.lock b/uv.lock index 8db25ae..e98519d 100644 --- a/uv.lock +++ b/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" diff --git a/view_summary.py b/view_summary.py index e60c73e..fd1c37b 100644 --- a/view_summary.py +++ b/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__":