From a6e54549e6b05fadd86bae6ae7e2e636a01d4df0 Mon Sep 17 00:00:00 2001 From: shijianglong Date: Wed, 10 Sep 2025 11:41:12 +0800 Subject: [PATCH 1/4] feat: add customizable memory_service extension for get_fast_api_app --- src/google/adk/cli/fast_api.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/google/adk/cli/fast_api.py b/src/google/adk/cli/fast_api.py index 7d93b54360..16357c1a93 100644 --- a/src/google/adk/cli/fast_api.py +++ b/src/google/adk/cli/fast_api.py @@ -60,6 +60,7 @@ def get_fast_api_app( session_db_kwargs: Optional[Mapping[str, Any]] = None, artifact_service_uri: Optional[str] = None, memory_service_uri: Optional[str] = None, + memory_service: Optional[Any] = None, eval_storage_uri: Optional[str] = None, allow_origins: Optional[list[str]] = None, web: bool, @@ -107,7 +108,9 @@ def _parse_agent_engine_resource_name(agent_engine_id_or_resource_name): return project, location, agent_engine_id # Build the Memory service - if memory_service_uri: + if memory_service: + pass + elif memory_service_uri: if memory_service_uri.startswith("rag://"): from ..memory.vertex_ai_rag_memory_service import VertexAiRagMemoryService From e34aec66d37df2e629cb580dc83862909e27ba27 Mon Sep 17 00:00:00 2001 From: shijianglong Date: Wed, 10 Sep 2025 12:44:13 +0800 Subject: [PATCH 2/4] fix: improve type hint for memory_service parameter --- src/google/adk/cli/fast_api.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/google/adk/cli/fast_api.py b/src/google/adk/cli/fast_api.py index 16357c1a93..7bef5317ac 100644 --- a/src/google/adk/cli/fast_api.py +++ b/src/google/adk/cli/fast_api.py @@ -38,6 +38,7 @@ from ..auth.credential_service.in_memory_credential_service import InMemoryCredentialService from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager +from ..memory.base_memory_service import BaseMemoryService from ..memory.in_memory_memory_service import InMemoryMemoryService from ..memory.vertex_ai_memory_bank_service import VertexAiMemoryBankService from ..runners import Runner @@ -60,7 +61,7 @@ def get_fast_api_app( session_db_kwargs: Optional[Mapping[str, Any]] = None, artifact_service_uri: Optional[str] = None, memory_service_uri: Optional[str] = None, - memory_service: Optional[Any] = None, + memory_service: Optional[BaseMemoryService] = None, eval_storage_uri: Optional[str] = None, allow_origins: Optional[list[str]] = None, web: bool, From 2337eb03b516b0d1cd72646473cbdb59145e2fa7 Mon Sep 17 00:00:00 2001 From: shijianglong Date: Thu, 11 Sep 2025 17:04:55 +0800 Subject: [PATCH 3/4] feature:add session end marker, when llm final response --- src/google/adk/cli/adk_web_server.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/google/adk/cli/adk_web_server.py b/src/google/adk/cli/adk_web_server.py index 4adce296a6..80b8de7c2a 100644 --- a/src/google/adk/cli/adk_web_server.py +++ b/src/google/adk/cli/adk_web_server.py @@ -1076,6 +1076,8 @@ async def event_generator(): "Generated event in agent run streaming: %s", sse_event ) yield f"data: {sse_event}\n\n" + # Add done flag when all data has been sent + yield "data: {\"done\": true}\n\n" except Exception as e: logger.exception("Error in event_generator: %s", e) # You might want to yield an error event here From add5f8c13f70fde9cafd99bd04a345d030a9bfc0 Mon Sep 17 00:00:00 2001 From: shijianglong Date: Thu, 11 Sep 2025 17:28:05 +0800 Subject: [PATCH 4/4] define SSE 'done' message as a constant --- src/google/adk/cli/adk_web_server.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/google/adk/cli/adk_web_server.py b/src/google/adk/cli/adk_web_server.py index 80b8de7c2a..61e2f7e4ba 100644 --- a/src/google/adk/cli/adk_web_server.py +++ b/src/google/adk/cli/adk_web_server.py @@ -1072,12 +1072,13 @@ async def event_generator(): sse_event = event.model_dump_json( exclude_none=True, by_alias=True ) + SSE_DONE_MESSAGE = 'data: {"done": true}\n\n' logger.debug( "Generated event in agent run streaming: %s", sse_event ) yield f"data: {sse_event}\n\n" # Add done flag when all data has been sent - yield "data: {\"done\": true}\n\n" + yield SSE_DONE_MESSAGE except Exception as e: logger.exception("Error in event_generator: %s", e) # You might want to yield an error event here