Skip to content

Commit d2d2231

Browse files
Formatting and fixing the bucket name handling
1 parent 96d1dca commit d2d2231

7 files changed

Lines changed: 244 additions & 51 deletions

File tree

contributing/samples/gepa/experiment.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,6 @@
4343
from tau_bench.types import EnvRunResult
4444
from tau_bench.types import RunConfig
4545
import tau_bench_agent as tau_bench_agent_lib
46-
4746
import utils
4847

4948

contributing/samples/gepa/run_experiment.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
from absl import flags
2626
import experiment
2727
from google.genai import types
28-
2928
import utils
3029

3130
_OUTPUT_DIR = flags.DEFINE_string(

src/google/adk/firestore_database_runner.py

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,9 @@
1414

1515
from __future__ import annotations
1616

17-
from typing import TYPE_CHECKING
17+
import os
1818
from typing import Optional
19+
from typing import TYPE_CHECKING
1920

2021
from .artifacts.gcs_artifact_service import GcsArtifactService
2122
from .memory.firestore_memory_service import FirestoreMemoryService
@@ -41,15 +42,13 @@ def create_firestore_runner(
4142
Returns:
4243
A Runner instance configured with Firestore services.
4344
"""
44-
# GcsArtifactService might require bucket name in constructor or read from env.
45-
# Let's assume it reads from env or takes it.
46-
# If we pass it, we might need to check its signature.
47-
# Let's assume it takes bucket_name if provided, or reads from env.
48-
artifact_service = GcsArtifactService()
49-
if gcs_bucket_name:
50-
# If GcsArtifactService supports setting it, we set it.
51-
# Or we can assume it reads from ADK_GCS_BUCKET_NAME env var.
52-
pass
45+
bucket_name = gcs_bucket_name or os.environ.get("ADK_GCS_BUCKET_NAME")
46+
if not bucket_name:
47+
raise ValueError(
48+
"Required property 'ADK_GCS_BUCKET_NAME' is not set. This"
49+
" is needed for the GcsArtifactService."
50+
)
51+
artifact_service = GcsArtifactService(bucket_name=bucket_name)
5352

5453
session_service = FirestoreSessionService(
5554
root_collection=firestore_root_collection

src/google/adk/memory/firestore_memory_service.py

Lines changed: 204 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@
2424
from google.cloud import firestore
2525
from typing_extensions import override
2626

27-
from ..events.event import Event
2827
from . import _utils
28+
from ..events.event import Event
2929
from .base_memory_service import BaseMemoryService
3030
from .base_memory_service import SearchMemoryResponse
3131
from .memory_entry import MemoryEntry
@@ -39,28 +39,206 @@
3939

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

66244

@@ -85,7 +263,9 @@ def __init__(
85263
"""
86264
self.client = client or firestore.AsyncClient()
87265
self.events_collection = events_collection or DEFAULT_EVENTS_COLLECTION
88-
self.stop_words = stop_words if stop_words is not None else DEFAULT_STOP_WORDS
266+
self.stop_words = (
267+
stop_words if stop_words is not None else DEFAULT_STOP_WORDS
268+
)
89269

90270
@override
91271
async def add_session_to_memory(self, session: Session) -> None:

src/google/adk/sessions/firestore_session_service.py

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
from .base_session_service import ListSessionsResponse
2929
from .session import Session
3030

31-
3231
logger = logging.getLogger("google_adk." + __name__)
3332

3433
DEFAULT_ROOT_COLLECTION = "adk-session"
@@ -65,7 +64,9 @@ def __init__(
6564
self.app_state_collection = DEFAULT_APP_STATE_COLLECTION
6665
self.user_state_collection = DEFAULT_USER_STATE_COLLECTION
6766

68-
def _get_sessions_ref(self, user_id: str) -> firestore.AsyncCollectionReference:
67+
def _get_sessions_ref(
68+
self, user_id: str
69+
) -> firestore.AsyncCollectionReference:
6970
return (
7071
self.client.collection(self.root_collection)
7172
.document(user_id)
@@ -83,6 +84,7 @@ async def create_session(
8384
"""Creates a new session in Firestore."""
8485
if not session_id:
8586
from google.adk.platform import uuid as platform_uuid
87+
8688
session_id = platform_uuid.new_uuid()
8789

8890
initial_state = state or {}
@@ -94,6 +96,7 @@ async def create_session(
9496
doc = await session_ref.get()
9597
if doc.exists:
9698
from ..errors.already_exists_error import AlreadyExistsError
99+
97100
raise AlreadyExistsError(f"Session {session_id} already exists.")
98101

99102
session_data = {
@@ -113,6 +116,7 @@ async def create_session(
113116
# the object, but the DB will have SERVER_TIMESTAMP.
114117
from datetime import datetime
115118
from datetime import timezone
119+
116120
local_now = datetime.now(timezone.utc).timestamp()
117121

118122
return Session(
@@ -323,13 +327,12 @@ async def append_event(self, session: Session, event: Event) -> Event:
323327

324328
for key, value in state_delta.items():
325329
if key.startswith("_app_"):
326-
app_updates[key[len("_app_"):]] = value
330+
app_updates[key[len("_app_") :]] = value
327331
elif key.startswith("_user_"):
328-
user_updates[key[len("_user_"):]] = value
332+
user_updates[key[len("_user_") :]] = value
329333
else:
330334
session_updates[key] = value
331335

332-
333336
# Update session doc with new state and updateTime
334337
# We'll do it outside the batch or inside if we can.
335338
# Let's use batch for everything to be atomic.
@@ -367,7 +370,9 @@ async def append_event(self, session: Session, event: Event) -> Event:
367370

368371
# Add event
369372
event_id = event.id
370-
event_ref = session_ref.collection(self.events_collection).document(event_id)
373+
event_ref = session_ref.collection(self.events_collection).document(
374+
event_id
375+
)
371376
# Store event data as JSON serialized string or dict
372377
event_data = event.model_dump(exclude_none=True, mode="json")
373378
batch.set(
@@ -385,7 +390,9 @@ async def append_event(self, session: Session, event: Event) -> Event:
385390
# No state delta, just add event and update session timestamp
386391
batch = self.client.batch()
387392
event_id = event.id
388-
event_ref = session_ref.collection(self.events_collection).document(event_id)
393+
event_ref = session_ref.collection(self.events_collection).document(
394+
event_id
395+
)
389396
event_data = event.model_dump(exclude_none=True, mode="json")
390397
batch.set(
391398
event_ref,

0 commit comments

Comments
 (0)