8000 feat: add _sync implementation in the inmemory session sevice. · DarioMR1/adk-python@bc43a11 · GitHub
[go: up one dir, main page]
Skip to content

Commit bc43a11

Browse files
DeanChensjcopybara-github
DeanChensj
authored and
copybara-github
committed
feat: add _sync implementation in the inmemory session sevice.
PiperOrigin-RevId: 758832846
1 parent dc90c91 commit bc43a11

File tree

1 file changed

+100
-2
lines changed

1 file changed

+100
-2
lines changed

‎src/google/adk/sessions/in_memory_session_service.py

Lines changed: 100 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414

1515
import copy
16+
import logging
1617
import time
1718
from typing import Any
1819
from typing import Optional
@@ -28,12 +29,15 @@
2829
from .session import Session
2930
from .state import State
3031

32+
logger = logging.getLogger(__name__)
33+
3134

3235
class InMemorySessionService(BaseSessionService):
3336
"""An in-memory implementation of the session service."""
3437

3538
def __init__(self):
36-
# A map from app name to a map from user ID to a map from session ID to session.
39+
# A map from app name to a map from user ID to a map from session ID to
40+
# session.
3741
self.sessions: dict[str, dict[str, dict[str, Session]]] = {}
3842
# A map from app name to a map from user ID to a map from key to the value.
3943
self.user_state: dict[str, dict[str, dict[str, Any]]] = {}
@@ -48,6 +52,37 @@ def create_session(
4852
user_id: str,
4953
state: Optional[dict[str, Any]] = None,
5054
session_id: Optional[str] = None,
55+
) -> Session:
56+
return self._create_session_impl(
57+
app_name=app_name,
58+
user_id=user_id,
59+
state=state,
60+
session_id=session_id,
61+
)
62+
63+
def create_session_sync(
64+
self,
65+
*,
66+
app_name: str,
67+
user_id: str,
68+
state: Optional[dict[str, Any]] = None,
69+
session_id: Optional[str] = None,
70+
) -> Session:
71+
logger.warning('Deprecated. Please migrate to the async method.')
72+
return self._create_session_impl(
73+
app_name=app_name,
74+
user_id=user_id,
75+
state=state,
76+
session_id=session_id,
77+
)
78+
79+
def _create_session_impl(
80+
self,
81+
*,
82+
app_name: str,
83+
user_id: str,
84+
state: Optional[dict[str, Any]] = None,
85+
session_id: Optional[str] = None,
5186
) -> Session:
5287
session_id = (
5388
session_id.strip()
@@ -79,6 +114,37 @@ def get_session(
79114
user_id: str,
80115
session_id: str,
81116
config: Optional[GetSessionConfig] = None,
117+
) -> Session:
118+
return self._get_session_impl(
119+
app_name=app_name,
120+
user_id=user_id,
121+
session_id=session_id,
122+
config=config,
123+
)
124+
125+
def get_session_sync(
126+
self,
127+
*,
128+
app_name: str,
129+
user_id: str,
130+
session_id: str,
131+
config: Optional[GetSessionConfig] = None,
132+
) -> Session:
133+
logger.warning('Deprecated. Please migrate to the async method.')
134+
return self._get_session_impl(
135+
app_name=app_name,
136+
user_id=user_id,
137+
session_id=session_id,
138+
config=config,
139+
)
140+
141+
def _get_session_impl(
142+
self,
143+
*,
144+
app_name: str,
145+
user_id: str,
146+
session_id: str,
147+
config: Optional[GetSessionConfig] = None,
82148
) -> Session:
83149
if app_name not in self.sessions:
84150
return None
@@ -130,6 +196,17 @@ def _merge_state(self, app_name: str, user_id: str, copied_session: Session):
130196
@override
131197
def list_sessions(
132198
self, *, app_name: str, user_id: str
199+
) -> ListSessionsResponse:
200+
return self._list_sessions_impl(app_name=app_name, user_id=user_id)
201+
202+
def list_sessions_sync(
203+
self, *, app_name: str, user_id: str
204+
) -> ListSessionsResponse:
205+
logger.warning('Deprecated. Please migrate to the async method.')
206+
return self._list_sessions_impl(app_name=app_name, user_id=user_id)
207+
208+
def _list_sessions_impl(
209+
self, *, app_name: str, user_id: str
133210
) -> ListSessionsResponse:
134211
empty_response = ListSessionsResponse()
135212
if app_name not in self.sessions:
@@ -145,9 +222,23 @@ def list_sessions(
145222
sessions_without_events.append(copied_session)
146223
return ListSessionsResponse(sessions=sessions_without_events)
147224

148-
@override
149225
def delete_session(
150226
self, *, app_name: str, user_id: str, session_id: str
227+
) -> None:
228+
self._delete_session_impl(
229+
app_name=app_name, user_id=user_id, session_id=session_id
230+
)
231+
232+
def delete_session_sync(
233+
self, *, app_name: str, user_id: str, session_id: str
234+
) -> None:
235+
logger.warning('Deprecated. Please migrate to the async method.')
236+
self._delete_session_impl(
237+
app_name=app_name, user_id=user_id, session_id=session_id
238+
)
239+
240+
def _delete_session_impl(
241+
self, *, app_name: str, user_id: str, session_id: str
151242
) -> None:
152243
if (
153244
self.get_session(
@@ -161,6 +252,13 @@ def delete_session(
161252

162253
@override
163254
def append_event(self, session: Session, event: Event) -> Event:
255+
return self._append_event_impl(session=session, event=event)
256+
257+
def append_event_sync(self, session: Session, event: Event) -> Event:
258+
logger.warning('Deprecated. Please migrate to the async method.')
259+
return self._append_event_impl(session=session, event=event)
260+
261+
def _append_event_impl(self, session: Session, event: Event) -> Event:
164262
# Update the in-memory session.
165263
super().append_event(session=session, event=event)
166264
session.last_update_time = event.timestamp

0 commit comments

Comments
 (0)
0