|
15 | 15 | from __future__ import annotations |
16 | 16 |
|
17 | 17 | import asyncio |
| 18 | +from contextlib import closing |
| 19 | +from datetime import datetime |
18 | 20 | from pathlib import Path |
19 | 21 | import sqlite3 |
| 22 | +from typing import Literal |
20 | 23 |
|
21 | 24 | import aiosqlite |
22 | 25 | from google.adk.events.event import Event |
@@ -351,6 +354,195 @@ async def test_persistence_across_restarts(tmp_path, fts_mode): |
351 | 354 | assert response.memories[0].custom_metadata["session_id"] == session.id |
352 | 355 |
|
353 | 356 |
|
| 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 | + |
354 | 546 | @pytest.mark.asyncio |
355 | 547 | async def test_reopen_database_with_fts_finds_earlier_rows_once( |
356 | 548 | tmp_path, require_fts5 |
|
0 commit comments