Spaces:
Sleeping
Sleeping
| # src/agent.py | |
| import os | |
| from pathlib import Path | |
| from google.adk.agents import LlmAgent | |
| from google.adk.agents.callback_context import CallbackContext | |
| from google.adk.models import LlmRequest | |
| from google.adk.models.google_llm import Gemini | |
| from google.adk.models.lite_llm import LiteLlm | |
| from google.adk.runners import Runner | |
| from google.adk.sessions import DatabaseSessionService, InMemorySessionService | |
| from google.genai import types | |
| from typing import Awaitable, Callable | |
| from src.config import settings | |
| from src.memory.core import get_block | |
| from src.memory.recall import SqliteRecallMemoryService | |
| from src.memory.archival import ArchivalStore | |
| from src.tools.send import send_message | |
| from src.tools.core_tools import core_memory_append, core_memory_replace | |
| from src.tools.search_tools import recall_search, archival_insert, archival_search | |
| APP_NAME = "memgpt_lite" | |
| DEFAULT_USER_ID = "default" | |
| _SYSTEM_PROMPT_PATH = Path(__file__).parent / "prompts" / "system.md" | |
| # Storage mode: "disk" (default, dev + demo recording) uses on-disk SQLite | |
| # + on-disk FAISS + DatabaseSessionService. "memory" (set on HF Space) uses | |
| # InMemorySessionService + :memory: SQLite + non-persisted FAISS — container | |
| # sleep naturally resets state, no cleanup code needed. See config.py for | |
| # MEMGPT_DATA_DIR (disk mode's data root). | |
| _STORAGE_MODE = os.environ.get("MEMGPT_STORAGE_MODE", "disk") | |
| if _STORAGE_MODE not in ("disk", "memory"): | |
| raise RuntimeError( | |
| f"MEMGPT_STORAGE_MODE must be 'disk' or 'memory', got {_STORAGE_MODE!r}" | |
| ) | |
| # Demo-abuse cap: bounds cost per LLM call. Default 1024 out; override with | |
| # MEMGPT_MAX_OUTPUT_TOKENS env var (2048 for demo recording; higher for | |
| # self-hosters). Documented in README §Deployment. | |
| _MAX_OUTPUT_TOKENS = int(os.environ.get("MEMGPT_MAX_OUTPUT_TOKENS", "1024")) | |
| if _STORAGE_MODE == "memory": | |
| session_service = InMemorySessionService() | |
| memory_service = SqliteRecallMemoryService(db_path=":memory:") | |
| archival_store = ArchivalStore( | |
| db_path=":memory:", | |
| index_path=None, | |
| dim=settings.embeddings.dim, | |
| ) | |
| else: | |
| session_service = DatabaseSessionService( | |
| db_url=f"sqlite+aiosqlite:///{Path(settings.storage.sessions_db).as_posix()}" | |
| ) | |
| memory_service = SqliteRecallMemoryService(db_path=settings.storage.recall_db) | |
| archival_store = ArchivalStore( | |
| db_path=settings.storage.archival_db, | |
| index_path=settings.storage.archival_index, | |
| dim=settings.embeddings.dim, | |
| ) | |
| def _build_model(): | |
| if settings.model.provider == "gemini": | |
| retry_config = types.HttpRetryOptions( | |
| attempts=settings.model.retry.attempts, | |
| exp_base=settings.model.retry.exp_base, | |
| initial_delay=settings.model.retry.initial_delay, | |
| http_status_codes=settings.model.retry.http_status_codes, | |
| ) | |
| return Gemini(model=settings.model.name, retry_options=retry_config) | |
| if settings.model.provider == "openrouter": | |
| return LiteLlm( | |
| model=settings.model.name, | |
| num_retries=settings.model.retry.attempts, | |
| base_url="https://openrouter.ai/api/v1", | |
| ) | |
| raise ValueError(f"Unknown model provider: {settings.model.provider}") | |
| def _inject_core_memory( | |
| callback_context: CallbackContext, llm_request: LlmRequest | |
| ) -> None: | |
| """Append <persona> and <human> blocks to system_instruction each turn.""" | |
| state = callback_context._invocation_context.session.state | |
| persona = get_block(state, "persona") | |
| human = get_block(state, "human") | |
| base = llm_request.config.system_instruction or "" | |
| if persona: | |
| base += f"\n\n<persona>\n{persona}\n</persona>" | |
| if human: | |
| base += f"\n\n<human>\n{human}\n</human>" | |
| llm_request.config.system_instruction = base | |
| async def _auto_save_to_memory(callback_context: CallbackContext) -> None: | |
| """Mirror every completed turn into recall memory.""" | |
| await callback_context._invocation_context.memory_service.add_session_to_memory( | |
| callback_context._invocation_context.session | |
| ) | |
| def build_agent() -> LlmAgent: | |
| return LlmAgent( | |
| model=_build_model(), | |
| name="memgpt_lite", | |
| description="Chatbot with tiered persistent memory (MemGPT-lite).", | |
| instruction=_SYSTEM_PROMPT_PATH.read_text(encoding="utf-8"), | |
| tools=[ | |
| send_message, | |
| core_memory_append, | |
| core_memory_replace, | |
| recall_search, | |
| archival_insert, | |
| archival_search, | |
| ], | |
| before_model_callback=_inject_core_memory, | |
| after_agent_callback=_auto_save_to_memory, | |
| # Demo-abuse cap: bounds cost per LLM call. Default 1024 out; override with | |
| # MEMGPT_MAX_OUTPUT_TOKENS env var (2048 for demo recording; | |
| # self-hosters can raise/remove). Documented in README §Deployment. | |
| generate_content_config=types.GenerateContentConfig( | |
| max_output_tokens=_MAX_OUTPUT_TOKENS, | |
| ), | |
| ) | |
| def build_runner() -> Runner: | |
| return Runner( | |
| agent=build_agent(), | |
| app_name=APP_NAME, | |
| session_service=session_service, | |
| memory_service=memory_service, | |
| ) | |
| async def _ensure_session(session_id: str, user_id: str): | |
| """Get an existing session or create a new one. | |
| Follows the notebook's try-create-except-get pattern: create_session | |
| raises if the id is taken; get_session returns the existing one. | |
| """ | |
| try: | |
| return await session_service.create_session( | |
| app_name=APP_NAME, user_id=user_id, session_id=session_id | |
| ) | |
| except Exception: | |
| return await session_service.get_session( | |
| app_name=APP_NAME, user_id=user_id, session_id=session_id | |
| ) | |
| ConfirmCallback = Callable[[str, dict], Awaitable[bool]] | |
| def _extract_reply(events) -> str: | |
| """Return the last send_message emission across the given events (may be '').""" | |
| reply = "" | |
| for event in events: | |
| if not event.content or not event.content.parts: | |
| continue | |
| for part in event.content.parts: | |
| fc = getattr(part, "function_call", None) | |
| if fc is not None and fc.name == "send_message": | |
| reply = fc.args.get("message", "") | |
| return reply | |
| async def start_turn( | |
| runner: Runner, | |
| user_text: str, | |
| session_id: str = "default", | |
| user_id: str = DEFAULT_USER_ID, | |
| ) -> dict: | |
| """Run phase 1 of a user turn. Returns a dict with reply + pending approval. | |
| Return shape: | |
| { | |
| "reply": str, # from send_message in phase-1 events (may be "") | |
| "pending": None | { | |
| "approval_id": str, | |
| "invocation_id": str, | |
| "hint": str, | |
| "payload": dict, | |
| "session_id": str, | |
| "user_id": str, | |
| } | |
| } | |
| If ``pending`` is ``None``, the turn completed. Otherwise the invocation | |
| is paused; the caller must invoke ``resume_turn(runner, pending, approved)`` | |
| to finish it. | |
| """ | |
| session = await _ensure_session(session_id, user_id) | |
| query = types.Content(role="user", parts=[types.Part(text=user_text)]) | |
| events = [] | |
| async for event in runner.run_async( | |
| user_id=user_id, session_id=session.id, new_message=query | |
| ): | |
| events.append(event) | |
| reply = _extract_reply(events) | |
| approval = _check_for_approval(events) | |
| if approval is None: | |
| return {"reply": reply, "pending": None} | |
| return { | |
| "reply": reply, | |
| "pending": { | |
| **approval, | |
| "session_id": session.id, | |
| "user_id": user_id, | |
| }, | |
| } | |
| async def resume_turn( | |
| runner: Runner, | |
| pending: dict, | |
| approved: bool, | |
| ) -> str: | |
| """Resume a paused invocation with an approval decision. Returns reply text.""" | |
| resume_msg = _build_approval_response(pending["approval_id"], approved) | |
| events = [] | |
| async for event in runner.run_async( | |
| user_id=pending["user_id"], | |
| session_id=pending["session_id"], | |
| new_message=resume_msg, | |
| invocation_id=pending["invocation_id"], | |
| ): | |
| events.append(event) | |
| return _extract_reply(events) | |
| async def run_turn( | |
| runner: Runner, | |
| user_text: str, | |
| session_id: str = "default", | |
| user_id: str = DEFAULT_USER_ID, | |
| on_confirm: ConfirmCallback | None = None, | |
| ) -> str: | |
| """Convenience wrapper: start_turn, auto-resume if paused. For REPL/tests. | |
| The Gradio UI in Step 18c uses start_turn + resume_turn directly instead, | |
| because it needs to hand control back to the user for the yes/no reply. | |
| """ | |
| result = await start_turn(runner, user_text, session_id, user_id) | |
| if result["pending"] is None: | |
| return result["reply"] | |
| approved = ( | |
| await on_confirm(result["pending"]["hint"], result["pending"]["payload"]) | |
| if on_confirm is not None | |
| else False | |
| ) | |
| resumed_reply = await resume_turn(runner, result["pending"], approved) | |
| return resumed_reply if resumed_reply else result["reply"] | |
| def _check_for_approval(events) -> dict | None: | |
| """Scan collected events for a confirmation request. Returns None if not paused. | |
| ADK nests the confirmation data under a ``toolConfirmation`` object inside | |
| ``fc.args`` (alongside ``originalFunctionCall``); the hint and payload are | |
| NOT at the top level. Inspected via scratch/inspect_confirmation_args.py — | |
| banking the shape here so a future ADK bump has a diff to grep for. | |
| """ | |
| for event in events: | |
| if not event.content or not event.content.parts: | |
| continue | |
| for part in event.content.parts: | |
| fc = getattr(part, "function_call", None) | |
| if fc is not None and fc.name == "adk_request_confirmation": | |
| confirmation = fc.args.get("toolConfirmation", {}) or {} | |
| return { | |
| "approval_id": fc.id, | |
| "invocation_id": event.invocation_id, | |
| "hint": confirmation.get("hint", ""), | |
| "payload": confirmation.get("payload", {}), | |
| } | |
| return | |
| def _build_approval_response(approval_id: str, approved: bool) -> types.Content: | |
| """Build the FunctionResponse Content that resumes a paused invocation.""" | |
| response = types.FunctionResponse( | |
| id=approval_id, | |
| name="adk_request_confirmation", | |
| response={"confirmed": approved}, | |
| ) | |
| return types.Content(role="user", parts=[types.Part(function_response=response)]) | |