feat: add endpoint to generate memory from session

Merge https://github.com/google/adk-python/pull/2900

In relation with #2416

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/2900 from guillaumeblaquiere:add-session-to-memory 0507de43021c62f9223167dca8f53b536227ad04
PiperOrigin-RevId: 808658162
This commit is contained in:
guillaume blaquiere
2025-09-18 11:13:21 -07:00
committed by Copybara-Service
parent 6b49391546
commit 25958242db
2 changed files with 58 additions and 1 deletions
+43
View File
@@ -31,6 +31,7 @@ from typing import Optional
from fastapi import FastAPI
from fastapi import HTTPException
from fastapi import Query
from fastapi import Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import RedirectResponse
from fastapi.responses import StreamingResponse
@@ -210,6 +211,13 @@ class RunEvalRequest(common.BaseModel):
eval_metrics: list[EvalMetric]
class UpdateMemoryRequest(common.BaseModel):
"""Request to add a session to the memory service."""
session_id: str
"""The ID of the session to add to memory."""
class RunEvalResult(common.BaseModel):
eval_set_file: str
eval_set_id: str
@@ -1144,6 +1152,41 @@ class AdkWebServer:
filename=artifact_name,
)
@app.patch("/apps/{app_name}/users/{user_id}/memory")
async def patch_memory(
app_name: str, user_id: str, update_memory_request: UpdateMemoryRequest
) -> None:
"""Adds all events from a given session to the memory service.
Args:
app_name: The name of the application.
user_id: The ID of the user.
update_memory_request: The memory request for the update
Raises:
HTTPException: If the memory service is not configured or the request is invalid.
"""
if not self.memory_service:
raise HTTPException(
status_code=400, detail="Memory service is not configured."
)
if (
update_memory_request is None
or update_memory_request.session_id is None
):
raise HTTPException(
status_code=400, detail="Update memory request is invalid."
)
session = await self.session_service.get_session(
app_name=app_name,
user_id=user_id,
session_id=update_memory_request.session_id,
)
if not session:
raise HTTPException(status_code=404, detail="Session not found")
await self.memory_service.add_session_to_memory(session)
@app.post("/run", response_model_exclude_none=True)
async def run_agent(req: RunAgentRequest) -> list[Event]:
session = await self.session_service.get_session(
+15 -1
View File
@@ -22,6 +22,7 @@ import tempfile
import time
from typing import Any
from typing import Optional
from unittest.mock import AsyncMock
from unittest.mock import MagicMock
from unittest.mock import patch
@@ -344,7 +345,7 @@ def mock_artifact_service():
@pytest.fixture
def mock_memory_service():
"""Create a mock memory service."""
return MagicMock()
return AsyncMock()
@pytest.fixture
@@ -939,5 +940,18 @@ def test_a2a_disabled_by_default(test_app):
logger.info("A2A disabled by default test passed")
def test_patch_memory(test_app, create_test_session, mock_memory_service):
"""Test adding a session to memory."""
info = create_test_session
url = f"/apps/{info['app_name']}/users/{info['user_id']}/memory"
payload = {"session_id": info["session_id"]}
response = test_app.patch(url, json=payload)
# Verify the response
assert response.status_code == 200
mock_memory_service.add_session_to_memory.assert_called_once()
logger.info("Add session to memory test completed successfully")
if __name__ == "__main__":
pytest.main(["-xvs", __file__])