Merge branch 'main' into fix-litellm-azure-pdf-upload

Author: SHAI-nikhil-chaudhary

src/google/adk/memory/in_memory_memory_service.py

@@ -32,6 +32,8 @@
 
 _UNKNOWN_SESSION_ID = '__unknown_session_id__'
 
+_MAX_SEARCH_RESULTS = 10
+
 
 def _user_key(app_name: str, user_id: str) -> tuple[str, str]:
   return (app_name, user_id)
@@ -45,7 +47,8 @@ def _extract_words_lower(text: str) -> set[str]:
 class InMemoryMemoryService(BaseMemoryService):
   """An in-memory memory service for prototyping purpose only.
 
-  Uses keyword matching instead of semantic search.
+  Uses keyword matching instead of semantic search. A search returns at most
+  ten memories, the ones sharing the most words with the query.
 
   This class is thread-safe, however, it should be used for testing and
   development only.
@@ -117,7 +120,7 @@ async def search_memory(
       ]
 
     words_in_query = _extract_words_lower(query)
-    response = SearchMemoryResponse()
+    scored_memories: list[tuple[int, MemoryEntry]] = []
 
     for session_events in session_event_lists:
       for event in session_events:
@@ -129,13 +132,23 @@ async def search_memory(
         if not words_in_event:
           continue
 
-        if any(query_word in words_in_event for query_word in words_in_query):
-          response.memories.append(
+        matched_words = len(words_in_query & words_in_event)
+        if matched_words:
+          scored_memories.append((
+              matched_words,
               MemoryEntry(
                   content=event.content,
                   author=event.author,
                   timestamp=_utils.format_timestamp(event.timestamp),
-              )
-          )
-
-    return response
+              ),
+          ))
+
+    # Almost any two sentences share a word, so returning every event that
+    # matches at least one query word returns most of the store, and callers
+    # such as the preload_memory tool put all of it in the prompt. Keep the
+    # events matching the most query words. The sort key reads only the count,
+    # so it is stable and events matching equally stay in insertion order.
+    scored_memories.sort(key=lambda scored_memory: -scored_memory[0])
+    return SearchMemoryResponse(
+        memories=[memory for _, memory in scored_memories[:_MAX_SEARCH_RESULTS]]
+    )

tests/unittests/isolated_import_utils.py

@@ -20,6 +20,7 @@
 
 from __future__ import annotations
 
+import json
 import os
 from pathlib import Path
 import subprocess
@@ -54,6 +55,33 @@ def run_isolated(source: str) -> subprocess.CompletedProcess[str]:
     )
 
 
+def loaded_top_level_packages(source: str) -> frozenset[str]:
+  """Returns the third-party top-level packages source leaves imported.
+
+  Standard-library modules, private modules and the pseudo-modules the
+  interpreter injects carry no install or startup cost of their own, so they
+  are dropped and only the distributions a caller pays for remain.
+  """
+  result = run_isolated(f"""
+import json
+import sys
+{source}
+
+names = {{
+    name.partition('.')[0]
+    for name, module in sys.modules.items()
+    if getattr(module, '__spec__', None) is not None
+}}
+print(json.dumps(sorted(
+    name
+    for name in names - sys.stdlib_module_names
+    if not name.startswith('_')
+)))
+""")
+  assert result.returncode == 0, result.stderr
+  return frozenset(json.loads(result.stdout.splitlines()[-1]))
+
+
 def assert_modules_unloaded(source: str, forbidden: tuple[str, ...]) -> None:
   """Asserts source leaves every forbidden module (and submodule) unimported."""
   result = run_isolated(f"""

tests/unittests/memory/test_in_memory_memory_service.py

@@ -393,6 +393,73 @@ async def test_search_memory_matches_non_latin_text():
   assert result.memories[0].content.parts[0].text == 'Привет мир'
 
 
+def _text_event(tag: str, text: str) -> Event:
+  return Event(
+      id=f'event-{tag}',
+      invocation_id=f'inv-{tag}',
+      author='user',
+      timestamp=1.0,
+      content=types.Content(parts=[types.Part(text=text)]),
+  )
+
+
+@pytest.mark.asyncio
+async def test_search_memory_ranks_by_number_of_matching_words():
+  """Tests that the events matching the most query words come first."""
+  memory_service = InMemoryMemoryService()
+  await memory_service.add_session_to_memory(
+      Session(
+          app_name=MOCK_APP_NAME,
+          user_id=MOCK_USER_ID,
+          id='session-ranked',
+          last_update_time=1000,
+          events=[
+              _text_event('ranked-a', 'The deploy is ready.'),
+              _text_event('ranked-b', 'Ready.'),
+              _text_event('ranked-c', 'The deploy status is ready.'),
+          ],
+      )
+  )
+
+  result = await memory_service.search_memory(
+      app_name=MOCK_APP_NAME, user_id=MOCK_USER_ID, query='deploy status ready'
+  )
+
+  assert [memory.content.parts[0].text for memory in result.memories] == [
+      'The deploy status is ready.',
+      'The deploy is ready.',
+      'Ready.',
+  ]
+
+
+@pytest.mark.asyncio
+async def test_search_memory_returns_at_most_ten_memories():
+  """Tests that a word shared with the whole store cannot return the store."""
+  memory_service = InMemoryMemoryService()
+  events = [_text_event(f'note-{i}', f'note {i} about work') for i in range(20)]
+  events.append(_text_event('backlog', 'the backlog note about work'))
+  await memory_service.add_session_to_memory(
+      Session(
+          app_name=MOCK_APP_NAME,
+          user_id=MOCK_USER_ID,
+          id='session-many',
+          last_update_time=1000,
+          events=events,
+      )
+  )
+
+  result = await memory_service.search_memory(
+      app_name=MOCK_APP_NAME, user_id=MOCK_USER_ID, query='work backlog note'
+  )
+
+  texts = [memory.content.parts[0].text for memory in result.memories]
+  # The best match is stored last but ranks first, and the rest tie, so they
+  # keep the order they were added in.
+  assert texts == ['the backlog note about work'] + [
+      f'note {i} about work' for i in range(9)
+  ]
+
+
 def _make_event(tag: str) -> Event:
   return Event(
       id=f'event-{tag}',

tests/unittests/test_import_loading.py

@@ -26,6 +26,7 @@
 
 from . import isolated_import_utils
 from .isolated_import_utils import assert_modules_unloaded
+from .isolated_import_utils import loaded_top_level_packages
 from .isolated_import_utils import run_isolated
 
 pytestmark = pytest.mark.skipif(
@@ -42,6 +43,59 @@
     'google.adk.workflow',
 )
 
+# The statements almost every ADK program starts with, and therefore the two
+# import graphs whose cost every user pays.
+_ENTRY_POINTS = (
+    'from google.adk.agents import Agent',
+    'from google.adk.runners import Runner',
+)
+
+# Third-party top-level packages an entry point may load. The forbidden lists
+# above pin individual deferrals on the lazy package inits; this one bounds the
+# whole graph, because the cost that reaches users arrives as a package nobody
+# noticed rather than as one somebody predicted.
+_ENTRY_POINT_PACKAGE_ALLOWLIST = frozenset({
+    # Declared requirements that ADK imports at module scope.
+    'click',
+    'fastapi',
+    'google',
+    'httpx',
+    'opentelemetry',
+    'packaging',
+    'pydantic',
+    'python_multipart',
+    'starlette',
+    'tenacity',
+    'websockets',
+    # Reached through pydantic and httpx rather than through ADK.
+    'annotated_doc',
+    'annotated_types',
+    'anyio',
+    'certifi',
+    'idna',
+    'orjson',
+    'pydantic_core',
+    'pygments',
+    'rich',
+    'sniffio',
+    'typing_extensions',
+    'typing_inspection',
+    'zstandard',
+    # google.genai.types annotates optional fields with aiohttp and Pillow
+    # types and imports whichever of the two the environment happens to have.
+    # No ADK module imports either one, so these are absent in some installs.
+    'PIL',
+    'aiohappyeyeballs',
+    'aiohttp',
+    'aiosignal',
+    'attr',
+    'defusedxml',
+    'frozenlist',
+    'multidict',
+    'propcache',
+    'yarl',
+})
+
 
 @pytest.mark.parametrize(
     ('module_name', 'forbidden'),
@@ -104,6 +158,29 @@ def test_package_import_defers_unrelated_runtime(
   )
 
 
+@pytest.mark.parametrize('statement', _ENTRY_POINTS, ids=('agent', 'runner'))
+def test_entry_point_loads_only_allowlisted_packages(statement: str) -> None:
+  """The two entry points every program uses load a reviewed set of packages.
+
+  The lazy package inits are already cheap, so a new eager dependency shows up
+  here first: as a package nobody agreed to pay for on every ADK start.
+
+  The unit is the top-level import name, so a new eager dependency arriving
+  under the `google` namespace, which ADK loads either way, does not show up
+  here.
+  """
+  unexpected = sorted(
+      loaded_top_level_packages(statement) - _ENTRY_POINT_PACKAGE_ALLOWLIST
+  )
+
+  assert not unexpected, (
+      f'{statement!r} now loads {", ".join(unexpected)}, which every ADK'
+      ' process would pay for at startup. Move the import into the function'
+      ' that needs it, or add the package to the allowlist together with the'
+      ' reason it has to be eager.'
+  )
+
+
 def test_constructing_agent_defers_optional_mcp_server_stack():
   """A normal Agent does not import MCP just because its extra is installed."""
   if importlib.util.find_spec('mcp') is None: