8000 ADK changes · GeekyprogrammerEJ/adk-python@bcf1deb · GitHub
[go: up one dir, main page]
Skip to content

Commit bcf1deb

Browse files
DeanChensjcopybara-github
DeanChensj
authored and
copybara-github
committed
ADK changes
PiperOrigin-RevId: 754131080
1 parent 8790643 commit bcf1deb

File tree

7 files changed

+15
-13
lines changed

7 files changed

+15
-13
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ class BaseMemoryService(abc.ABC):
5151
"""
5252

5353
@abc.abstractmethod
54-
def add_session_to_memory(self, session: Session):
54+
async def add_session_to_memory(self, session: Session):
5555
"""Adds a session to the memory service.
5656
5757
A session may be added multiple times during its lifetime.
@@ -61,7 +61,7 @@ def add_session_to_memory(self, session: Session):
6161
"""
6262

6363
@abc.abstractmethod
64-
def search_memory(
64+
async def search_memory(
6565
self, *, app_name: str, user_id: str, query: str
6666
) -> SearchMemoryResponse:
6767
"""Searches for sessions that match the query.

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,13 @@ def __init__(self):
2929
self.session_events: dict[str, list[Event]] = {}
3030
"""keys are app_name/user_id/session_id"""
3131

32-
def add_session_to_memory(self, session: Session):
32+
async def add_session_to_memory(self, session: Session):
3333
key = f'{session.app_name}/{session.user_id}/{session.id}'
3434
self.session_events[key] = [
3535
event for event in session.events if event.content
3636
]
3737

38-
def search_memory(
38+
async def search_memory(
3939
self, *, app_name: str, user_id: str, query: str
4040
) -> SearchMemoryResponse:
4141
"""Prototyping purpose only."""

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ def __init__(
5454
)
5555

5656
@override
57-
def add_session_to_memory(self, session: Session):
57+
async def add_session_to_memory(self, session: Session):
5858
with tempfile.NamedTemporaryFile(
5959
mode="w", delete=False, suffix=".txt"
6060
) as temp_file:
@@ -91,7 +91,7 @@ def add_session_to_memory(self, session: Session):
9191
os.remove(temp_file_path)
9292

9393
@override
94 8000 -
def search_memory(
94+
async def search_memory(
9595
self, *, app_name: str, user_id: str, query: str
9696
) -> SearchMemoryResponse:
9797
"""Searches for sessions that match the query using rag.retrieval_query."""

‎src/google/adk/runners.py

Lines changed: 2 additions & 2 deletions
+
await self.memory_service.add_session_to_memory(session)
Original file line numberDiff line numberDiff line change
@@ -297,14 +297,14 @@ async def run_live(
297297
self.session_service.append_event(session=session, event=event)
298298
yield event
299299

300-
def close_session(self, session: Session):
300+
async def close_session(self, session: Session):
301301
"""Closes a session and adds it to the memory service (experimental feature).
302302
303303
Args:
304304
session: The session to close.
305305
"""
306306
if self.memory_service:
307-
self.memory_service.add_session_to_memory(session)
307
308308
self.session_service.close_session(session=session)
309309

310310
def _find_agent_to_run(

‎src/google/adk/tools/load_memory_tool.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,9 @@
2727
from ..models import LlmRequest
2828

2929

30-
def load_memory(query: str, tool_context: ToolContext) -> 'list[MemoryResult]':
30+
async def load_memory(
31+
query: str, tool_context: ToolContext
32+
) -> 'list[MemoryResult]':
3133
"""Loads the memory for the current user.
3234
3335
Args:
@@ -36,7 +38,7 @@ def load_memory(query: str, tool_context: ToolContext) -> 'list[MemoryResult]':
3638
Returns:
3739
A list of memory results.
3840
"""
39-
response = tool_context.search_memory(query)
41+
response = await tool_context.search_memory(query)
4042
return response.memories
4143

4244

‎src/google/adk/tools/preload_memory_tool.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ async def process_llm_request(
4545
if not parts or not parts[0].text:
4646
return
4747
query = parts[0].text
48-
response = tool_context.search_memory(query)
48+
response = await tool_context.search_memory(query)
4949
if not response.memories:
5050
return
5151
memory_text = ''

‎src/google/adk/tools/tool_context.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -79,11 +79,11 @@ def list_artifacts(self) -> list[str]:
7979
session_id=self._invocation_context.session.id,
8080
)
8181

82-
def search_memory(self, query: str) -> 'SearchMemoryResponse':
82+
async def search_memory(self, query: str) -> SearchMemoryResponse:
8383
"""Searches the memory of the current user."""
8484
if self._invocation_context.memory_service is None:
8585
raise ValueError('Memory service is not available.')
86-
return self._invocation_context.memory_service.search_memory(
86+
return await self._invocation_context.memory_service.search_memory(
8787
app_name=self._invocation_context.app_name,
8888
user_id=self._invocation_context.user_id,
8989
query=query,

0 commit comments

Comments
 (0)
0