Skip to content

Commit 28a571e

Browse files
ScottMansfieldScottMansfield
committed
pyink and isort
1 parent bf47b5c commit 28a571e

5 files changed

Lines changed: 248 additions & 43 deletions

File tree

‎src/google/adk/integrations/firestore/firestore_memory_service.py‎

Lines changed: 201 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@
2424

2525
from typing_extensions import override
2626

27-
from ...memory import _utils
2827
from ...events.event import Event
28+
from ...memory import _utils
2929
from ...memory.base_memory_service import BaseMemoryService
3030
from ...memory.base_memory_service import SearchMemoryResponse
3131
from ...memory.memory_entry import MemoryEntry
@@ -40,28 +40,206 @@
4040
DEFAULT_EVENTS_COLLECTION = "events"
4141

4242
DEFAULT_STOP_WORDS = {
43-
"a", "an", "the", "and", "or", "but", "if", "then", "else", "to", "of",
44-
"in", "on", "for", "with", "is", "are", "was", "were", "be", "been",
45-
"being", "have", "has", "had", "do", "does", "did", "can", "could",
46-
"will", "would", "should", "shall", "may", "might", "must", "up", "down",
47-
"out", "in", "over", "under", "again", "further", "then", "once", "here",
48-
"there", "when", "where", "why", "how", "all", "any", "both", "each",
49-
"few", "more", "most", "other", "some", "such", "no", "nor", "not",
50-
"only", "own", "same", "so", "than", "too", "very", "i", "me", "my",
51-
"myself", "we", "our", "ours", "ourselves", "you", "your", "yours",
52-
"yourself", "yourselves", "he", "him", "his", "himself", "she", "her",
53-
"hers", "herself", "it", "its", "itself", "they", "them", "their",
54-
"theirs", "themselves", "what", "which", "who", "whom", "this", "that",
55-
"these", "those", "am", "is", "are", "was", "were", "be", "been", "being",
56-
"have", "has", "had", "having", "do", "does", "did", "doing", "a", "an",
57-
"the", "and", "but", "if", "or", "because", "as", "until", "while", "of",
58-
"at", "by", "for", "with", "about", "against", "between", "into",
59-
"through", "during", "before", "after", "above", "below", "to", "from",
60-
"up", "down", "in", "out", "on", "off", "over", "under", "again",
61-
"further", "then", "once", "here", "there", "when", "where", "why", "how",
62-
"all", "any", "both", "each", "few", "more", "most", "other", "some",
63-
"such", "no", "nor", "not", "only", "own", "same", "so", "than", "too",
64-
"very", "s", "t", "can", "will", "just", "don", "should", "now",
43+
"a",
44+
"an",
45+
"the",
46+
"and",
47+
"or",
48+
"but",
49+
"if",
50+
"then",
51+
"else",
52+
"to",
53+
"of",
54+
"in",
55+
"on",
56+
"for",
57+
"with",
58+
"is",
59+
"are",
60+
"was",
61+
"were",
62+
"be",
63+
"been",
64+
"being",
65+
"have",
66+
"has",
67+
"had",
68+
"do",
69+
"does",
70+
"did",
71+
"can",
72+
"could",
73+
"will",
74+
"would",
75+
"should",
76+
"shall",
77+
"may",
78+
"might",
79+
"must",
80+
"up",
81+
"down",
82+
"out",
83+
"in",
84+
"over",
85+
"under",
86+
"again",
87+
"further",
88+
"then",
89+
"once",
90+
"here",
91+
"there",
92+
"when",
93+
"where",
94+
"why",
95+
"how",
96+
"all",
97+
"any",
98+
"both",
99+
"each",
100+
"few",
101+
"more",
102+
"most",
103+
"other",
104+
"some",
105+
"such",
106+
"no",
107+
"nor",
108+
"not",
109+
"only",
110+
"own",
111+
"same",
112+
"so",
113+
"than",
114+
"too",
115+
"very",
116+
"i",
117+
"me",
118+
"my",
119+
"myself",
120+
"we",
121+
"our",
122+
"ours",
123+
"ourselves",
124+
"you",
125+
"your",
126+
"yours",
127+
"yourself",
128+
"yourselves",
129+
"he",
130+
"him",
131+
"his",
132+
"himself",
133+
"she",
134+
"her",
135+
"hers",
136+
"herself",
137+
"it",
138+
"its",
139+
"itself",
140+
"they",
141+
"them",
142+
"their",
143+
"theirs",
144+
"themselves",
145+
"what",
146+
"which",
147+
"who",
148+
"whom",
149+
"this",
150+
"that",
151+
"these",
152+
"those",
153+
"am",
154+
"is",
155+
"are",
156+
"was",
157+
"were",
158+
"be",
159+
"been",
160+
"being",
161+
"have",
162+
"has",
163+
"had",
164+
"having",
165+
"do",
166+
"does",
167+
"did",
168+
"doing",
169+
"a",
170+
"an",
171+
"the",
172+
"and",
173+
"but",
174+
"if",
175+
"or",
176+
"because",
177+
"as",
178+
"until",
179+
"while",
180+
"of",
181+
"at",
182+
"by",
183+
"for",
184+
"with",
185+
"about",
186+
"against",
187+
"between",
188+
"into",
189+
"through",
190+
"during",
191+
"before",
192+
"after",
193+
"above",
194+
"below",
195+
"to",
196+
"from",
197+
"up",
198+
"down",
199+
"in",
200+
"out",
201+
"on",
202+
"off",
203+
"over",
204+
"under",
205+
"again",
206+
"further",
207+
"then",
208+
"once",
209+
"here",
210+
"there",
211+
"when",
212+
"where",
213+
"why",
214+
"how",
215+
"all",
216+
"any",
217+
"both",
218+
"each",
219+
"few",
220+
"more",
221+
"most",
222+
"other",
223+
"some",
224+
"such",
225+
"no",
226+
"nor",
227+
"not",
228+
"only",
229+
"own",
230+
"same",
231+
"so",
232+
"than",
233+
"too",
234+
"very",
235+
"s",
236+
"t",
237+
"can",
238+
"will",
239+
"just",
240+
"don",
241+
"should",
242+
"now",
65243
}
66244

67245

‎src/google/adk/integrations/firestore/firestore_session_service.py‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ def __init__(
6464

6565
self.client = client or firestore.AsyncClient()
6666
self.root_collection = (
67-
root_collection
67+
root_collection
6868
or os.environ.get("ADK_FIRESTORE_ROOT_COLLECTION")
6969
or DEFAULT_ROOT_COLLECTION
7070
)
@@ -108,6 +108,7 @@ async def create_session(
108108
) -> Session:
109109
"""Creates a new session in Firestore."""
110110
from google.cloud import firestore
111+
111112
if not session_id:
112113
from ...platform import uuid as platform_uuid
113114

@@ -283,7 +284,7 @@ async def delete_session(
283284
session_ref = self._get_sessions_ref(user_id).document(session_id)
284285

285286
events_ref = session_ref.collection(self.events_collection)
286-
287+
287288
batch = self.client.batch()
288289
count = 0
289290
async for event_doc in events_ref.stream():
@@ -303,7 +304,10 @@ async def _update_app_state_transactional(
303304
) -> dict[str, Any]:
304305
"""Atomically applies delta to app state inside a transaction."""
305306
from google.cloud import firestore
306-
doc_ref = self.client.collection(self.app_state_collection).document(app_name)
307+
308+
doc_ref = self.client.collection(self.app_state_collection).document(
309+
app_name
310+
)
307311

308312
@firestore.async_transactional
309313
async def _txn(transaction):
@@ -321,6 +325,7 @@ async def _update_user_state_transactional(
321325
) -> dict[str, Any]:
322326
"""Atomically applies delta to user state inside a transaction."""
323327
from google.cloud import firestore
328+
324329
doc_ref = (
325330
self.client.collection(self.user_state_collection)
326331
.document(app_name)
@@ -342,6 +347,7 @@ async def _txn(transaction):
342347
async def append_event(self, session: Session, event: Event) -> Event:
343348
"""Appends an event to a session in Firestore."""
344349
from google.cloud import firestore
350+
345351
if event.partial:
346352
return event
347353

@@ -365,10 +371,14 @@ async def append_event(self, session: Session, event: Event) -> Event:
365371
session_updates[key] = value
366372

367373
if app_updates:
368-
await self._update_app_state_transactional(session.app_name, app_updates)
374+
await self._update_app_state_transactional(
375+
session.app_name, app_updates
376+
)
369377

370378
if user_updates:
371-
await self._update_user_state_transactional(session.app_name, session.user_id, user_updates)
379+
await self._update_user_state_transactional(
380+
session.app_name, session.user_id, user_updates
381+
)
372382

373383
for k, v in session_updates.items():
374384
session.state[k] = v

‎tests/unittests/integrations/firestore/test_firestore_database_runner.py‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,9 @@ def test_create_firestore_runner_with_arg(mock_agent, monkeypatch):
3535
mock.patch(
3636
"google.adk.integrations.firestore.firestore_database_runner.FirestoreSessionService"
3737
),
38-
mock.patch("google.adk.integrations.firestore.firestore_database_runner.FirestoreMemoryService"),
38+
mock.patch(
39+
"google.adk.integrations.firestore.firestore_database_runner.FirestoreMemoryService"
40+
),
3941
mock.patch(
4042
"google.adk.integrations.firestore.firestore_database_runner.GcsArtifactService"
4143
) as mock_gcs,
@@ -53,7 +55,9 @@ def test_create_firestore_runner_with_env(mock_agent, monkeypatch):
5355
mock.patch(
5456
"google.adk.integrations.firestore.firestore_database_runner.FirestoreSessionService"
5557
),
56-
mock.patch("google.adk.integrations.firestore.firestore_database_runner.FirestoreMemoryService"),
58+
mock.patch(
59+
"google.adk.integrations.firestore.firestore_database_runner.FirestoreMemoryService"
60+
),
5761
mock.patch(
5862
"google.adk.integrations.firestore.firestore_database_runner.GcsArtifactService"
5963
) as mock_gcs,
@@ -80,13 +84,16 @@ def test_create_firestore_runner_with_root_collection(mock_agent, monkeypatch):
8084
mock.patch(
8185
"google.adk.integrations.firestore.firestore_database_runner.FirestoreSessionService"
8286
) as mock_session,
83-
mock.patch("google.adk.integrations.firestore.firestore_database_runner.FirestoreMemoryService"),
84-
mock.patch("google.adk.integrations.firestore.firestore_database_runner.GcsArtifactService"),
87+
mock.patch(
88+
"google.adk.integrations.firestore.firestore_database_runner.FirestoreMemoryService"
89+
),
90+
mock.patch(
91+
"google.adk.integrations.firestore.firestore_database_runner.GcsArtifactService"
92+
),
8593
):
8694
runner = create_firestore_runner(
8795
mock_agent, firestore_root_collection="custom_collection"
8896
)
8997

9098
assert runner is not None
9199
mock_session.assert_called_once_with(root_collection="custom_collection")
92-

‎tests/unittests/integrations/firestore/test_firestore_memory_service.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,9 @@ async def test_search_memory_deduplication(mock_firestore_client):
127127

128128
get_mock = mock.AsyncMock(side_effect=[[doc_snapshot1], [doc_snapshot2]])
129129

130-
mock_firestore_client.collection_group.return_value.where.return_value.where.return_value.where.return_value.get = get_mock
130+
mock_firestore_client.collection_group.return_value.where.return_value.where.return_value.where.return_value.get = (
131+
get_mock
132+
)
131133

132134
response = await service.search_memory(
133135
app_name=app_name, user_id=user_id, query=query
@@ -145,7 +147,9 @@ async def test_search_memory_parsing_error(mock_firestore_client, caplog):
145147
user_id = "test_user"
146148
query = "quick"
147149

148-
doc_snapshot = mock_firestore_client.collection_group.return_value.where.return_value.where.return_value.where.return_value.get.return_value[0]
150+
doc_snapshot = mock_firestore_client.collection_group.return_value.where.return_value.where.return_value.where.return_value.get.return_value[
151+
0
152+
]
149153
doc_snapshot.to_dict.return_value = {"event_data": "invalid_data"}
150154

151155
response = await service.search_memory(
@@ -164,4 +168,3 @@ async def test_search_memory_only_stop_words(mock_firestore_client):
164168
)
165169
assert not response.memories
166170
mock_firestore_client.collection_group.assert_not_called()
167-

0 commit comments

Comments
 (0)