Merge pull request 'refactor: timestamp compatible with Firestore' (#30) from refactor/timestamp-to-date into main
Some checks failed
CI / ci (push) Failing after 12s
Some checks failed
CI / ci (push) Failing after 12s
Reviewed-on: #30
This commit was merged in pull request #30.
This commit is contained in:
@@ -6,6 +6,7 @@ import asyncio
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING, Any, override
|
||||
|
||||
from google.adk.errors.already_exists_error import AlreadyExistsError
|
||||
@@ -102,6 +103,24 @@ class FirestoreSessionService(BaseSessionService):
|
||||
def _events_col(self, app_name: str, user_id: str, session_id: str) -> Any:
|
||||
return self._session_ref(app_name, user_id, session_id).collection("events")
|
||||
|
||||
@staticmethod
|
||||
def _timestamp_to_float(value: Any, default: float = 0.0) -> float:
|
||||
if value is None:
|
||||
return default
|
||||
if isinstance(value, (int, float)):
|
||||
return float(value)
|
||||
if hasattr(value, "timestamp"):
|
||||
try:
|
||||
return float(value.timestamp())
|
||||
except (
|
||||
TypeError,
|
||||
ValueError,
|
||||
OSError,
|
||||
OverflowError,
|
||||
) as exc: # pragma: no cover
|
||||
logger.debug("Failed to convert timestamp %r: %s", value, exc)
|
||||
return default
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# State helpers
|
||||
# ------------------------------------------------------------------
|
||||
@@ -171,7 +190,7 @@ class FirestoreSessionService(BaseSessionService):
|
||||
)
|
||||
)
|
||||
|
||||
now = time.time()
|
||||
now = datetime.now(UTC)
|
||||
write_coros.append(
|
||||
self._session_ref(app_name, user_id, session_id).set(
|
||||
{
|
||||
@@ -196,7 +215,7 @@ class FirestoreSessionService(BaseSessionService):
|
||||
user_id=user_id,
|
||||
id=session_id,
|
||||
state=merged,
|
||||
last_update_time=now,
|
||||
last_update_time=now.timestamp(),
|
||||
)
|
||||
|
||||
@override
|
||||
@@ -283,7 +302,9 @@ class FirestoreSessionService(BaseSessionService):
|
||||
id=session_id,
|
||||
state=merged,
|
||||
events=events,
|
||||
last_update_time=session_data.get("last_update_time", 0.0),
|
||||
last_update_time=self._timestamp_to_float(
|
||||
session_data.get("last_update_time"), 0.0
|
||||
),
|
||||
)
|
||||
|
||||
@override
|
||||
@@ -326,7 +347,9 @@ class FirestoreSessionService(BaseSessionService):
|
||||
id=data["session_id"],
|
||||
state=merged,
|
||||
events=[],
|
||||
last_update_time=data.get("last_update_time", 0.0),
|
||||
last_update_time=self._timestamp_to_float(
|
||||
data.get("last_update_time"), 0.0
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
@@ -366,6 +389,8 @@ class FirestoreSessionService(BaseSessionService):
|
||||
# Persist state deltas
|
||||
session_ref = self._session_ref(app_name, user_id, session_id)
|
||||
|
||||
last_update_dt = datetime.fromtimestamp(event.timestamp, UTC)
|
||||
|
||||
if event.actions and event.actions.state_delta:
|
||||
state_deltas = _session_util.extract_state_delta(event.actions.state_delta)
|
||||
|
||||
@@ -386,16 +411,16 @@ class FirestoreSessionService(BaseSessionService):
|
||||
FieldPath("state", k).to_api_repr(): v
|
||||
for k, v in state_deltas["session"].items()
|
||||
}
|
||||
field_updates["last_update_time"] = event.timestamp
|
||||
field_updates["last_update_time"] = last_update_dt
|
||||
write_coros.append(session_ref.update(field_updates))
|
||||
else:
|
||||
write_coros.append(
|
||||
session_ref.update({"last_update_time": event.timestamp})
|
||||
session_ref.update({"last_update_time": last_update_dt})
|
||||
)
|
||||
|
||||
await asyncio.gather(*write_coros)
|
||||
else:
|
||||
await session_ref.update({"last_update_time": event.timestamp})
|
||||
await session_ref.update({"last_update_time": last_update_dt})
|
||||
|
||||
# Log token usage
|
||||
if event.usage_metadata:
|
||||
|
||||
Reference in New Issue
Block a user