Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 16 additions & 15 deletions src/google/adk/memory/vertex_ai_rag_memory_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,9 @@ def __init__(

@override
async def add_session_to_memory(self, session: Session) -> None:
if not self._vertex_rag_store.rag_resources:
raise ValueError("Rag resources must be set.")

with tempfile.NamedTemporaryFile(
mode="w", delete=False, suffix=".txt"
) as temp_file:
Expand All @@ -150,23 +153,21 @@ async def add_session_to_memory(self, session: Session) -> None:
temp_file.write(output_string)
temp_file_path = temp_file.name

if not self._vertex_rag_store.rag_resources:
raise ValueError("Rag resources must be set.")

from ..dependencies.vertexai import rag

for rag_resource in self._vertex_rag_store.rag_resources:
rag.upload_file(
corpus_name=rag_resource.rag_corpus,
path=temp_file_path,
# this is the temp workaround as upload file does not support
# adding metadata, thus use display_name to store the session info.
display_name=_build_source_display_name(
session.app_name, session.user_id, session.id
),
)

os.remove(temp_file_path)
try:
for rag_resource in self._vertex_rag_store.rag_resources:
rag.upload_file(
corpus_name=rag_resource.rag_corpus,
path=temp_file_path,
# this is the temp workaround as upload file does not support
# adding metadata, thus use display_name to store the session info.
display_name=_build_source_display_name(
session.app_name, session.user_id, session.id
),
)
finally:
os.remove(temp_file_path)

@override
async def search_memory(
Expand Down
68 changes: 68 additions & 0 deletions tests/unittests/memory/test_vertex_ai_rag_memory_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@
# limitations under the License.

import json
import os
import tempfile
from types import SimpleNamespace

from google.adk.events.event import Event
Expand All @@ -31,6 +33,72 @@ def _rag_context(source_display_name: str, text: str) -> SimpleNamespace:
)


def _sample_session() -> Session:
return Session(
app_name="demo",
user_id="alice",
id="session-1",
last_update_time=1,
events=[
Event(
id="event-1",
author="user",
timestamp=1,
content=types.Content(parts=[types.Part(text="hello")]),
)
],
)


def _track_temp_files(mocker) -> list[str]:
"""Records the paths of every NamedTemporaryFile created."""
created_paths: list[str] = []
real_named_temp_file = tempfile.NamedTemporaryFile

def _spy(*args, **kwargs):
temp_file = real_named_temp_file(*args, **kwargs)
created_paths.append(temp_file.name)
return temp_file

mocker.patch(
"google.adk.memory.vertex_ai_rag_memory_service.tempfile.NamedTemporaryFile",
side_effect=_spy,
)
return created_paths


@pytest.mark.asyncio
async def test_add_session_removes_temp_file_when_upload_fails(mocker):
"""The temp file must not leak when rag.upload_file raises."""
memory_service = VertexAiRagMemoryService(rag_corpus="unused")
created_paths = _track_temp_files(mocker)
fake_rag = SimpleNamespace(
upload_file=mocker.Mock(side_effect=RuntimeError("upload boom"))
)
mocker.patch("google.adk.dependencies.vertexai.rag", fake_rag)

with pytest.raises(RuntimeError, match="upload boom"):
await memory_service.add_session_to_memory(_sample_session())

assert created_paths, "expected a temp file to have been created"
assert not os.path.exists(created_paths[0])


@pytest.mark.asyncio
async def test_add_session_does_not_create_temp_file_without_rag_resources(
mocker,
):
"""Validation happens before the temp file is created, so nothing leaks."""
memory_service = VertexAiRagMemoryService(rag_corpus="unused")
memory_service._vertex_rag_store.rag_resources = None
created_paths = _track_temp_files(mocker)

with pytest.raises(ValueError, match="Rag resources must be set."):
await memory_service.add_session_to_memory(_sample_session())

assert created_paths == []


@pytest.mark.asyncio
async def test_search_memory_rejects_ambiguous_legacy_display_names(mocker):
"""Ensures dotted user IDs cannot match another user's legacy memory."""
Expand Down