Skip to content

Commit de34d8d

Browse files
xxwtiancaicopybara-github
authored andcommitted
fix: exclude thought parts from SQLite memory search results
Merge #7691 SqliteMemoryService.search_memory now drops parts marked thought=True from the content of each returned MemoryEntry, matching InMemoryMemoryService. Stored rows, including ones written before this change, are not rewritten. Fixes #7688 PiperOrigin-RevId: 1007534085
1 parent 799fc94 commit de34d8d

3 files changed

Lines changed: 272 additions & 0 deletions

File tree

‎src/google/adk/memory/_sqlite_memory_service.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -553,6 +553,8 @@ def _escape_like_query(word: str) -> str:
553553

554554
def _row_to_memory_entry(row: aiosqlite.Row) -> MemoryEntry:
555555
content = types.Content.model_validate_json(row["content_json"])
556+
if content.parts:
557+
content.parts = [part for part in content.parts if not part.thought]
556558
return MemoryEntry(
557559
id=str(row["id"]),
558560
content=content,

‎tests/unittests/memory/test_sqlite_memory_service.py‎

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,11 @@
1515
from __future__ import annotations
1616

1717
import asyncio
18+
from contextlib import closing
19+
from datetime import datetime
1820
from pathlib import Path
1921
import sqlite3
22+
from typing import Literal
2023

2124
import aiosqlite
2225
from google.adk.events.event import Event
@@ -351,6 +354,195 @@ async def test_persistence_across_restarts(tmp_path, fts_mode):
351354
assert response.memories[0].custom_metadata["session_id"] == session.id
352355

353356

357+
@pytest.mark.parametrize("in_memory", [False, True])
358+
@pytest.mark.parametrize("ingestion", ["events", "session"])
359+
async def test_search_returns_visible_parts_without_mutating_events(
360+
tmp_path: Path,
361+
fts_mode: Literal["off", "on"],
362+
in_memory: bool,
363+
ingestion: str,
364+
) -> None:
365+
"""Search removes thoughts while preserving visible content and identity."""
366+
db_path = tmp_path / "memory.db"
367+
service = SqliteMemoryService(
368+
":memory:" if in_memory else db_path, fts=fts_mode
369+
)
370+
visible_parts = [
371+
types.Part(text="My preferred tea is saffron.", thought=False),
372+
types.Part.from_bytes(data=b"visible image", mime_type="image/png"),
373+
types.Part(text="Serve it hot."),
374+
]
375+
event = Event(
376+
id="event-1",
377+
author="assistant",
378+
timestamp=1.0,
379+
content=types.Content(
380+
role="model",
381+
parts=[
382+
types.Part(text="Synthetic obsidiansecret", thought=True),
383+
visible_parts[0],
384+
types.Part(
385+
inline_data=types.Blob(
386+
data=b"synthetic thought image", mime_type="image/png"
387+
),
388+
thought=True,
389+
),
390+
*visible_parts[1:],
391+
],
392+
),
393+
)
394+
original_event = event.model_copy(deep=True)
395+
try:
396+
if ingestion == "session":
397+
await service.add_session_to_memory(_make_session([event]))
398+
else:
399+
await service.add_events_to_memory(
400+
app_name="app",
401+
user_id="user",
402+
session_id="session-1",
403+
events=[event],
404+
)
405+
if not in_memory:
406+
with closing(sqlite3.connect(db_path)) as conn:
407+
stored_before_search = conn.execute(
408+
"SELECT content_json FROM memory_events"
409+
).fetchone()[0]
410+
411+
response = await service.search_memory(
412+
app_name="app", user_id="user", query="saffron"
413+
)
414+
415+
if not in_memory:
416+
with closing(sqlite3.connect(db_path)) as conn:
417+
stored_after_search = conn.execute(
418+
"SELECT content_json FROM memory_events"
419+
).fetchone()[0]
420+
assert stored_before_search == stored_after_search
421+
assert stored_after_search == original_event.content.model_dump_json(
422+
exclude_none=True
423+
)
424+
assert event == original_event
425+
assert len(response.memories) == 1
426+
memory = response.memories[0]
427+
assert memory.content == types.Content(role="model", parts=visible_parts)
428+
assert memory.id == "1"
429+
assert memory.author == "assistant"
430+
assert memory.timestamp == datetime.fromtimestamp(1.0).isoformat()
431+
assert memory.custom_metadata == {
432+
"session_id": "session-1",
433+
"event_id": "event-1",
434+
}
435+
finally:
436+
await service.close()
437+
438+
439+
async def test_search_does_not_match_thought_text(
440+
fts_mode: Literal["off", "on"],
441+
) -> None:
442+
"""Thought text is not searchable even when the event has visible text."""
443+
service = SqliteMemoryService(":memory:", fts=fts_mode)
444+
event = Event(
445+
author="assistant",
446+
content=types.Content(
447+
role="model",
448+
parts=[
449+
types.Part(text="Synthetic obsidiansecret", thought=True),
450+
types.Part(text="My preferred tea is saffron."),
451+
],
452+
),
453+
)
454+
try:
455+
await service.add_events_to_memory(
456+
app_name="app", user_id="user", events=[event]
457+
)
458+
459+
response = await service.search_memory(
460+
app_name="app", user_id="user", query="obsidiansecret"
461+
)
462+
463+
assert not response.memories
464+
finally:
465+
await service.close()
466+
467+
468+
async def test_search_filters_preexisting_thoughts_without_rewriting_storage(
469+
tmp_path: Path, fts_mode: Literal["off", "on"]
470+
) -> None:
471+
"""Reopening old rows filters returned thoughts without rewriting JSON.
472+
473+
Seed stored JSON directly so this also catches an ingestion-only fix.
474+
"""
475+
db_path = tmp_path / "memory.db"
476+
service = SqliteMemoryService(db_path, fts=fts_mode)
477+
try:
478+
await service.search_memory(app_name="app", user_id="user", query="saffron")
479+
finally:
480+
await service.close()
481+
stored_content = """{
482+
"role": "model",
483+
"parts": [
484+
{"text": "Synthetic obsidiansecret", "thought": true},
485+
{"text": "My preferred tea is saffron.", "thought": false},
486+
{"file_data": {"file_uri": "gs://test-bucket/image.png",
487+
"mime_type": "image/png"}}
488+
]
489+
}"""
490+
with closing(sqlite3.connect(db_path)) as conn:
491+
conn.execute(
492+
"""INSERT INTO memory_events (
493+
id, app_name, user_id, session_id, event_id, author, timestamp,
494+
content_json, search_text
495+
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
496+
(
497+
42,
498+
"app",
499+
"user",
500+
"old-session",
501+
"old-event",
502+
"assistant",
503+
1.0,
504+
stored_content,
505+
"my preferred tea is saffron.",
506+
),
507+
)
508+
conn.commit()
509+
content_before_search = conn.execute(
510+
"SELECT content_json FROM memory_events WHERE id = 42"
511+
).fetchone()[0]
512+
513+
reopened_service = SqliteMemoryService(db_path, fts=fts_mode)
514+
try:
515+
response = await reopened_service.search_memory(
516+
app_name="app", user_id="user", query="saffron"
517+
)
518+
finally:
519+
await reopened_service.close()
520+
with closing(sqlite3.connect(db_path)) as conn:
521+
content_after_search = conn.execute(
522+
"SELECT content_json FROM memory_events WHERE id = 42"
523+
).fetchone()[0]
524+
525+
assert content_before_search == content_after_search == stored_content
526+
assert len(response.memories) == 1
527+
memory = response.memories[0]
528+
assert memory.content == types.Content(
529+
role="model",
530+
parts=[
531+
types.Part(text="My preferred tea is saffron.", thought=False),
532+
types.Part.from_uri(
533+
file_uri="gs://test-bucket/image.png", mime_type="image/png"
534+
),
535+
],
536+
)
537+
assert memory.id == "42"
538+
assert memory.author == "assistant"
539+
assert memory.timestamp == datetime.fromtimestamp(1.0).isoformat()
540+
assert memory.custom_metadata == {
541+
"session_id": "old-session",
542+
"event_id": "old-event",
543+
}
544+
545+
354546
@pytest.mark.asyncio
355547
async def test_reopen_database_with_fts_finds_earlier_rows_once(
356548
tmp_path, require_fts5

‎tests/unittests/tools/test_load_memory_tool.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,18 +12,96 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
from contextlib import closing
16+
import sqlite3
17+
from typing import Literal
1518
from unittest import mock
1619
from unittest.mock import AsyncMock
1720

21+
from google.adk.agents.llm_agent import Agent
22+
from google.adk.events.event import Event
1823
from google.adk.features import FeatureName
1924
from google.adk.features._feature_registry import temporary_feature_override
25+
from google.adk.memory import SqliteMemoryService
2026
from google.adk.models.llm_request import LlmRequest
2127
from google.adk.tools.load_memory_tool import load_memory_tool
2228
from google.adk.tools.preload_memory_tool import PreloadMemoryTool
2329
from google.adk.tools.tool_context import ToolContext
2430
from google.genai import types
2531
import pytest
2632

33+
from .. import testing_utils
34+
35+
36+
@pytest.mark.parametrize('fts_mode', ['off', 'on'])
37+
async def test_runner_loads_only_visible_sqlite_memory(
38+
fts_mode: Literal['off', 'on'],
39+
) -> None:
40+
"""The next model request contains visible memory without thought parts."""
41+
if fts_mode == 'on':
42+
with closing(sqlite3.connect(':memory:')) as conn:
43+
try:
44+
conn.execute('CREATE VIRTUAL TABLE test_fts USING fts5(col)')
45+
except sqlite3.OperationalError:
46+
pytest.skip('FTS5 not available in this SQLite build.')
47+
memory_service = SqliteMemoryService(':memory:', fts=fts_mode)
48+
model = testing_utils.MockModel.create(
49+
responses=[
50+
types.Part.from_function_call(
51+
name='load_memory', args={'query': 'saffron'}
52+
),
53+
types.Part(text='Your preferred tea is saffron.'),
54+
]
55+
)
56+
runner = testing_utils.InMemoryRunner(
57+
Agent(name='assistant', model=model, tools=[load_memory_tool])
58+
)
59+
runner.runner.memory_service = memory_service
60+
try:
61+
await memory_service.add_events_to_memory(
62+
app_name=runner.app_name,
63+
user_id='test_user',
64+
session_id='history',
65+
events=[
66+
Event(
67+
author='assistant',
68+
content=types.Content(
69+
role='model',
70+
parts=[
71+
types.Part(
72+
text='Synthetic obsidiansecret', thought=True
73+
),
74+
types.Part(text='My preferred tea is saffron.'),
75+
],
76+
),
77+
)
78+
],
79+
)
80+
81+
await runner.run_async('What is my preferred tea?')
82+
83+
assert len(model.requests) == 2
84+
function_responses = [
85+
part.function_response
86+
for content in model.requests[1].contents
87+
for part in content.parts or []
88+
if part.function_response
89+
and part.function_response.name == 'load_memory'
90+
]
91+
assert len(function_responses) == 1
92+
payload = function_responses[0].model_dump(mode='json')
93+
memories = payload['response']['result']['memories']
94+
assert len(memories) == 1
95+
assert types.Content.model_validate(
96+
memories[0]['content']
97+
) == types.Content(
98+
role='model', parts=[types.Part(text='My preferred tea is saffron.')]
99+
)
100+
assert 'obsidiansecret' not in model.requests[1].model_dump_json()
101+
finally:
102+
await memory_service.close()
103+
await runner.runner.close()
104+
27105

28106
def test_get_declaration_with_json_schema_feature_disabled():
29107
"""Test that _get_declaration uses parameters when feature is disabled."""

0 commit comments

Comments
 (0)